Big (2x - 5x) perf gain via momoisation

Author: AlexeyRaga

src/Hedgehog.Experimental.Tests/Hedgehog.Experimental.Tests.fsproj

@@ -5,7 +5,7 @@
     <OutputType>Exe</OutputType>
     <GenerateProgramFile>false</GenerateProgramFile>
   </PropertyGroup>
-
+  
   <ItemGroup>
     <Compile Include="ComplexGenericTest.fs" />
     <Compile Include="TypeParamMappingTests.fs" />

src/Hedgehog.Experimental/GeneratorCollection.fs

@@ -1,6 +1,7 @@
 namespace Hedgehog
 
 open System
+open System.Collections.Concurrent
 open System.Collections.Immutable
 
 // A generator factory which can be backed by a generic method.
@@ -21,8 +22,7 @@ type internal GeneratorKey = {
   ConcreteTypes: Type option list
 }
 
-[<Struct>]
-type internal GeneratorCollection =
+type internal GeneratorCollection = {
   // A dictionary of generators.
   // The key distinguishes between different patterns of generic vs concrete type arguments
   // The value is a tuple of:
@@ -31,14 +31,17 @@ type internal GeneratorCollection =
   // 3. A generator factory, which can be backed by a generic method,
   //    so it takes an array of genetic type parameters,
   //    and an array of arguments to create the generator.
-  GeneratorCollection of ImmutableDictionary<GeneratorKey, Type * Type[] * GeneratorFactory>
+  Generators: ImmutableDictionary<GeneratorKey, Type * Type[] * GeneratorFactory>
+  // The cache memoizes tryFindFor results to avoid repeated reflection calls.
+  Cache: ConcurrentDictionary<Type, (Type * (Type[] * GeneratorFactory)) option>
+}
 
 module internal GeneratorCollection =
 
-  let empty = GeneratorCollection(ImmutableDictionary.Empty)
-
-  let unwrap (GeneratorCollection map) = map
-  let map f = unwrap >> f >> GeneratorCollection
+  let empty = {
+    Generators = ImmutableDictionary.Empty
+    Cache = ConcurrentDictionary<Type, (Type * (Type[] * GeneratorFactory)) option>()
+  }
 
   /// Create a GeneratorKey from a type by identifying which positions are generic vs concrete
   let private createKey (t: Type) : GeneratorKey =
@@ -52,12 +55,21 @@ module internal GeneratorCollection =
       // Non-generic types use themselves as the key
       { GenericTypeDefinition = t; ConcreteTypes = [] }
 
-  let merge (GeneratorCollection gens1) (GeneratorCollection gens2) =
-    GeneratorCollection (gens1.SetItems(gens2))
+  let merge (gens1: GeneratorCollection) (gens2: GeneratorCollection) =
+    // Create a new cache when merging since the generator set has changed
+    {
+      Generators = gens1.Generators.SetItems(gens2.Generators)
+      Cache = ConcurrentDictionary<Type, (Type * (Type[] * GeneratorFactory)) option>()
+    }
 
   let addGenerator (normalizedType: Type) (originalType: Type) (paramTypes: Type[]) (factory: Type[] -> obj[] -> obj) =
-    let key = createKey normalizedType
-    map _.SetItem(key, (originalType, paramTypes, factory))
+    fun (gc: GeneratorCollection) ->
+      let key = createKey normalizedType
+      // Reset cache when adding a generator since lookup results may change
+      {
+        Generators = gc.Generators.SetItem(key, (originalType, paramTypes, factory))
+        Cache = ConcurrentDictionary<Type, (Type * (Type[] * GeneratorFactory)) option>()
+      }
 
   /// Count the number of generic parameters in a type
   let private countGenericParameters (t: Type) =
@@ -69,20 +81,23 @@ module internal GeneratorCollection =
   // to satisfy specific types (like Either<int, string>).
   // When multiple generators match, returns the most specific one (fewest generic parameters).
   // Returns the original reflected type along with the args and factory.
-  let tryFindFor (targetType: Type) =
-    unwrap
-    >> Seq.choose (fun (KeyValue (key, (originalType, paramTypes, factory))) ->
-        // Only consider generators with the same generic type definition
-        let targetKey = createKey targetType
-        if key.GenericTypeDefinition = targetKey.GenericTypeDefinition then
-          // Check if the stored type can satisfy the target type
-          if originalType |> TypeUtils.satisfies targetType then
-            Some (originalType, paramTypes, factory)
+  // Results are memoized to avoid repeated reflection calls.
+  let tryFindFor (targetType: Type) (gc: GeneratorCollection) =
+    gc.Cache.GetOrAdd(targetType, fun targetType ->
+      let targetKey = createKey targetType
+      gc.Generators
+      |> Seq.choose (fun (KeyValue (key, (originalType, paramTypes, factory))) ->
+          // Only consider generators with the same generic type definition
+          if key.GenericTypeDefinition = targetKey.GenericTypeDefinition then
+            // Check if the stored type can satisfy the target type
+            if originalType |> TypeUtils.satisfies targetType then
+              Some (originalType, paramTypes, factory)
+            else
+              None
           else
             None
-        else
-          None
+      )
+      |> Seq.sortBy (fun (originalType, _, _) -> countGenericParameters originalType)
+      |> Seq.tryHead
+      |> Option.map (fun (originalType, paramTypes, factory) -> (originalType, (paramTypes, factory)))
     )
-    >> Seq.sortBy (fun (originalType, _, _) -> countGenericParameters originalType)
-    >> Seq.tryHead
-    >> Option.map (fun (originalType, paramTypes, factory) -> (originalType, (paramTypes, factory)))