From 1d83525acde604937eb2843eb898af969c0c1ca3 Mon Sep 17 00:00:00 2001
From: David Sagan <dcs16@cornell.edu>
Date: Fri, 17 Jul 2026 14:30:37 -0400
Subject: [PATCH] Reduce load-time invalidations from untyped operator methods
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The out-of-place `*` for TaylorMaps and the `SymplecticS` `+/-/*//`
operators were defined with an untyped first (or second) argument, e.g.
`*(m2, m1::DAMap)` and `*(M, S::SymplecticS)`. Inserting these
`f(::Any, ...)` methods at load time invalidates large amounts of
precompiled generic code that calls the same functions.

Constrain the arguments to the types actually accepted:
- `*(m2, m1::TaylorMap)` is split into `*(m2::TaylorMap, m1::TaylorMap)`
  (map composition) and `*(m2::Union{Number,AbstractArray}, m1::TaylorMap)`
  (TPS scalar/vector function composition).
- `∘(m2, m1::TaylorMap)` gains the same `Union{Number,AbstractArray}`
  constraint on `m2`.
- The `SymplecticS` operators constrain the matrix argument to
  `AbstractVecOrMat` (they already index it via `size(M, .)`).

Measured with SnoopCompile on `using SciBmad`: NonlinearNormalForm's
invalidation children drop from 197 to 56 (the removed 141 were exactly
the untyped `*`/`SymplecticS` roots). No behavior change for valid
inputs; full test suite passes.

Refs bmad-sim/SciBmad.jl#77

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 src/map.jl   | 13 +++++++++----
 src/utils.jl |  4 ++--
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/map.jl b/src/map.jl
index fab3c77..85d7736 100644
--- a/src/map.jl
+++ b/src/map.jl
@@ -731,8 +731,10 @@ function ∘(m2::$t, m1::$t)
   return m
 end
 
-# When composing a TPS scalar/vector function w a map, use orbital part of map:
-function ∘(m2, m1::$t)
+# When composing a TPS scalar/vector function w a map, use orbital part of map.
+# `m2` is constrained to Number/AbstractArray (a TPS scalar or vector function) so
+# that inserting this method does not invalidate compiled `∘(::Any, ...)` callers.
+function ∘(m2::Union{Number,AbstractArray}, m1::$t)
   TI.is_tps_type(eltype(m2)) isa TI.IsTPSType || error("Cannot compose: $(eltype(m2)) is not a TPS type supported by TPSAInterface.jl")
   T = promote_type(eltype(m1.v), eltype(m2))
   T == eltype(m1.v) ? m1xprom = m1.v : m1xprom = T.(m1.v)
@@ -751,8 +753,11 @@ literal_pow(::typeof(^), m::$t{V0,V,Q,S}, vn::Val{-1}) where {V0,V,Q,S} = inv(m)
 inv(m::$t; do_spin::Bool=true) = (out_m = zero(m); inv!(out_m, m, do_spin=do_spin); return out_m)
 ^(m::$t, n::Integer) = (out_m = zero(m); pow!(out_m, m, n); return out_m)
 
-# Also allow * for simplicity and \ and / 
-*(m2, m1::$t) = ∘(m2, m1)
+# Also allow * for simplicity and \ and /
+# Split the former untyped `*(m2, m1::$t)` into concrete-argument methods so that
+# inserting them does not invalidate compiled `*(::Any, ...)` callers.
+*(m2::$t, m1::$t) = ∘(m2, m1)
+*(m2::Union{Number,AbstractArray}, m1::$t) = ∘(m2, m1)
 /(m2::$t, m1::$t) = m2 ∘ inv(m1) 
 \(m2::$t, m1::$t) = inv(m2) ∘ m1
 
diff --git a/src/utils.jl b/src/utils.jl
index 53203e1..1b10bf7 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -475,8 +475,8 @@ const S = SymplecticS()
 
 for op = (:+, :-, :*, :/)
 @eval begin
-Base.$op(S::SymplecticS,M) = Base.$op(JMatrix{Int8,+1}(size(M,1)), M)
-Base.$op(M,S::SymplecticS) = Base.$op(M, JMatrix{Int8,+1}(size(M,2)))
+Base.$op(S::SymplecticS,M::AbstractVecOrMat) = Base.$op(JMatrix{Int8,+1}(size(M,1)), M)
+Base.$op(M::AbstractVecOrMat,S::SymplecticS) = Base.$op(M, JMatrix{Int8,+1}(size(M,2)))
 end
 end