Move Moshi to extension for 17.5% load time reduction

Project.toml

@@ -18,14 +18,11 @@ IteratorInterfaceExtensions = "82899510-4779-5014-852e-03e436cf321d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
-Moshi = "2e0e35c7-a2e4-4343-998d-7ef72827ed2d"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
-RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
-RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
@@ -37,20 +34,26 @@ ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+Moshi = "2e0e35c7-a2e4-4343-998d-7ef72827ed2d"
 PartialFunctions = "570af359-4316-4cb7-8c74-252c00c2016b"
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 RCall = "6f49c342-dc21-5d91-9882-a32aef131414"
+RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
+RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 SciMLBaseChainRulesCoreExt = "ChainRulesCore"
 SciMLBaseMLStyleExt = "MLStyle"
 SciMLBaseMakieExt = "Makie"
+SciMLBaseMoshiExt = "Moshi"
 SciMLBasePartialFunctionsExt = "PartialFunctions"
 SciMLBasePyCallExt = "PyCall"
 SciMLBasePythonCallExt = "PythonCall"
 SciMLBaseRCallExt = "RCall"
+SciMLBaseRecipesBaseExt = "RecipesBase"
+SciMLBaseRuntimeGeneratedFunctionsExt = "RuntimeGeneratedFunctions"
 SciMLBaseZygoteExt = ["Zygote", "ChainRulesCore"]
 
 [compat]
@@ -73,7 +76,7 @@ Logging = "1.10"
 MLStyle = "0.4.17"
 Makie = "0.20, 0.21, 0.22, 0.23, 0.24"
 Markdown = "1.10"
-Moshi = "0.3"
+Moshi = "0.3.7"
 PartialFunctions = "1.1"
 PrecompileTools = "1.2"
 Preferences = "1.3"

ext/SciMLBaseMoshiExt.jl

@@ -0,0 +1,82 @@
+module SciMLBaseMoshiExt
+
+if isdefined(Base, :get_extension)
+    using Moshi.Data: @data
+    using Moshi.Match: @match
+else
+    using ..Moshi.Data: @data
+    using ..Moshi.Match: @match
+end
+
+using SciMLBase
+using SciMLBase: AbstractTimeseriesSolution, TimeDomain, PeriodicClock, SolverStepClock,
+                 ContinuousClock
+
+# Enhanced clock predicates using @match when Moshi is available
+# These override the fallback implementations with pattern matching
+function SciMLBase.isclock_moshi(c::TimeDomain)
+    @match c begin
+        PeriodicClock() => true
+        _ => false
+    end
+end
+
+function SciMLBase.issolverstepclock_moshi(c::TimeDomain)
+    @match c begin
+        SolverStepClock() => true
+        _ => false
+    end
+end
+
+function SciMLBase.iscontinuous_moshi(c::TimeDomain)
+    @match c begin
+        ContinuousClock() => true
+        _ => false
+    end
+end
+
+function SciMLBase.first_clock_tick_time_moshi(c, t0)
+    @match c begin
+        PeriodicClock(dt) => ceil(t0 / dt) * dt
+        SolverStepClock() => t0
+        ContinuousClock() => error("ContinuousClock() is not a discrete clock")
+    end
+end
+
+function SciMLBase.canonicalize_indexed_clock_moshi(ic::SciMLBase.IndexedClock, sol::AbstractTimeseriesSolution)
+    c = ic.clock
+
+    return @match c begin
+        PeriodicClock(dt) => ceil(sol.prob.tspan[1] / dt) * dt .+ (ic.idx .- 1) .* dt
+        SolverStepClock() => begin
+            ssc_idx = findfirst(eachindex(sol.discretes)) do i
+                !isa(sol.discretes[i].t, AbstractRange)
+            end
+            sol.discretes[ssc_idx].t[ic.idx]
+        end
+        ContinuousClock() => sol.t[ic.idx]
+    end
+end
+
+# Override fallback implementations to use Moshi versions
+function SciMLBase.isclock(c::TimeDomain)
+    SciMLBase.isclock_moshi(c)
+end
+
+function SciMLBase.issolverstepclock(c::TimeDomain)
+    SciMLBase.issolverstepclock_moshi(c)
+end
+
+function SciMLBase.iscontinuous(c::TimeDomain)
+    SciMLBase.iscontinuous_moshi(c)
+end
+
+function SciMLBase.first_clock_tick_time(c, t0)
+    SciMLBase.first_clock_tick_time_moshi(c, t0)
+end
+
+function SciMLBase.canonicalize_indexed_clock(ic::SciMLBase.IndexedClock, sol::AbstractTimeseriesSolution)
+    SciMLBase.canonicalize_indexed_clock_moshi(ic, sol)
+end
+
+end # module