Merge pull request #5345 from unisonweb/perf-bugfix

fix bindNames performance by building namespace+locals names only once
unisonweb · Sep 14, 2024 · 5bccedb · 5bccedb
2 parents eda2f0e + f1ca9d4
commit 5bccedb
Show file tree

Hide file tree

Showing 3 changed files with 120 additions and 104 deletions.
diff --git a/unison-core/src/Unison/Names.hs b/unison-core/src/Unison/Names.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE MultiWayIf #-}
 {-# LANGUAGE RecordWildCards #-}
 
 module Unison.Names
@@ -507,27 +508,28 @@ lenientToNametree names =
 -- determine what the name resolves to, per the usual suffix-matching rules (where local defnintions and direct
 -- dependencies are preferred to indirect dependencies).
 resolveName :: forall ref. (Ord ref) => Relation Name ref -> Set Name -> Name -> Set (ResolvesTo ref)
-resolveName namespace locals name
-  | Set.member name locals = Set.singleton (ResolvesToLocal name)
-  | Set.size exactNamespaceMatches == 1 = Set.mapMonotonic ResolvesToNamespace exactNamespaceMatches
-  | otherwise = localsPlusNamespaceSuffixMatches
+resolveName namespace locals =
+  \name ->
+    let exactNamespaceMatches :: Set ref
+        exactNamespaceMatches =
+          Relation.lookupDom name namespace
+        localsPlusNamespaceSuffixMatches :: Set (ResolvesTo ref)
+        localsPlusNamespaceSuffixMatches =
+          Name.searchByRankedSuffix name localsPlusNamespace
+     in if
+          | Set.member name locals -> Set.singleton (ResolvesToLocal name)
+          | Set.size exactNamespaceMatches == 1 -> Set.mapMonotonic ResolvesToNamespace exactNamespaceMatches
+          | otherwise -> localsPlusNamespaceSuffixMatches
   where
-    exactNamespaceMatches :: Set ref
-    exactNamespaceMatches =
-      Relation.lookupDom name namespace
-
-    localsPlusNamespaceSuffixMatches :: Set (ResolvesTo ref)
-    localsPlusNamespaceSuffixMatches =
-      Name.searchByRankedSuffix
-        name
-        ( shadowing1
-            ( List.foldl'
-                (\acc name -> Relation.insert name (ResolvesToLocal name) acc)
-                Relation.empty
-                (Set.toList locals)
-            )
-            ( Relation.map
-                (over _2 ResolvesToNamespace)
-                namespace
-            )
+    localsPlusNamespace :: Relation Name (ResolvesTo ref)
+    localsPlusNamespace =
+      shadowing1
+        ( List.foldl'
+            (\acc name -> Relation.insert name (ResolvesToLocal name) acc)
+            Relation.empty
+            (Set.toList locals)
+        )
+        ( Relation.map
+            (over _2 ResolvesToNamespace)
+            namespace
         )
diff --git a/unison-core/src/Unison/Term.hs b/unison-core/src/Unison/Term.hs
@@ -156,44 +156,51 @@ bindNames ::
   Names ->
   Term v a ->
   Names.ResolutionResult a (Term v a)
-bindNames unsafeVarToName nameToVar localVars namespace term = do
-  let freeTmVars = ABT.freeVarOccurrences localVars term
-      freeTyVars =
-        [ (v, a) | (v, as) <- Map.toList (freeTypeVarAnnotations term), a <- as
-        ]
-
-      okTm :: (v, a) -> Maybe (v, ResolvesTo Referent)
-      okTm (v, _) =
-        case Set.size matches of
-          1 -> Just (v, Set.findMin matches)
-          0 -> Nothing -- not found: leave free for telling user about expected type
-          _ -> Nothing -- ambiguous: leave free for TDNR
-        where
-          matches :: Set (ResolvesTo Referent)
-          matches =
-            Names.resolveName (Names.terms namespace) (Set.map unsafeVarToName localVars) (unsafeVarToName v)
-
-      okTy :: (v, a) -> Names.ResolutionResult a (v, Type v a)
-      okTy (v, a) =
-        case Names.lookupHQType Names.IncludeSuffixes hqName namespace of
-          rs
-            | Set.size rs == 1 -> pure (v, Type.ref a $ Set.findMin rs)
-            | Set.size rs == 0 -> Left (Seq.singleton (Names.TypeResolutionFailure hqName a Names.NotFound))
-            | otherwise -> Left (Seq.singleton (Names.TypeResolutionFailure hqName a (Names.Ambiguous namespace rs Set.empty)))
-        where
-          hqName = HQ.NameOnly (unsafeVarToName v)
-
-  let (namespaceTermResolutions, localTermResolutions) =
-        partitionResolutions (mapMaybe okTm freeTmVars)
-
-      termSubsts =
-        [(v, fromReferent () ref) | (v, ref) <- namespaceTermResolutions]
-          ++ [(v, var () (nameToVar name)) | (v, name) <- localTermResolutions]
-  typeSubsts <- validate okTy freeTyVars
-  pure $
-    term
-      & ABT.substsInheritAnnotation termSubsts
-      & substTypeVars typeSubsts
+bindNames unsafeVarToName nameToVar localVars namespace =
+  -- term is bound here because the where-clause binds a data structure that we only want to compute once, then share
+  -- across all calls to `bindNames` with different terms
+  \term -> do
+    let freeTmVars = ABT.freeVarOccurrences localVars term
+        freeTyVars =
+          [ (v, a) | (v, as) <- Map.toList (freeTypeVarAnnotations term), a <- as
+          ]
+
+        okTm :: (v, a) -> Maybe (v, ResolvesTo Referent)
+        okTm (v, _) =
+          case Set.size matches of
+            1 -> Just (v, Set.findMin matches)
+            0 -> Nothing -- not found: leave free for telling user about expected type
+            _ -> Nothing -- ambiguous: leave free for TDNR
+          where
+            matches :: Set (ResolvesTo Referent)
+            matches =
+              resolveTermName (unsafeVarToName v)
+
+        okTy :: (v, a) -> Names.ResolutionResult a (v, Type v a)
+        okTy (v, a) =
+          case Names.lookupHQType Names.IncludeSuffixes hqName namespace of
+            rs
+              | Set.size rs == 1 -> pure (v, Type.ref a $ Set.findMin rs)
+              | Set.size rs == 0 -> Left (Seq.singleton (Names.TypeResolutionFailure hqName a Names.NotFound))
+              | otherwise -> Left (Seq.singleton (Names.TypeResolutionFailure hqName a (Names.Ambiguous namespace rs Set.empty)))
+          where
+            hqName = HQ.NameOnly (unsafeVarToName v)
+
+    let (namespaceTermResolutions, localTermResolutions) =
+          partitionResolutions (mapMaybe okTm freeTmVars)
+
+        termSubsts =
+          [(v, fromReferent () ref) | (v, ref) <- namespaceTermResolutions]
+            ++ [(v, var () (nameToVar name)) | (v, name) <- localTermResolutions]
+    typeSubsts <- validate okTy freeTyVars
+    pure $
+      term
+        & ABT.substsInheritAnnotation termSubsts
+        & substTypeVars typeSubsts
+  where
+    resolveTermName :: Name.Name -> Set (ResolvesTo Referent)
+    resolveTermName =
+      Names.resolveName (Names.terms namespace) (Set.map unsafeVarToName localVars)
 
 -- Prepare a term for type-directed name resolution by replacing
 -- any remaining free variables with blanks to be resolved by TDNR

diff --git a/unison-core/src/Unison/Type/Names.hs b/unison-core/src/Unison/Type/Names.hs
@@ -28,49 +28,56 @@ bindNames ::
   Names ->
   Type v a ->
   Names.ResolutionResult a (Type v a)
-bindNames unsafeVarToName nameToVar localVars namespace ty =
-  let -- Identify the unresolved variables in the type: those whose names aren't an *exact* match for some locally-bound
-      -- type.
-      --
-      -- For example:
-      --
-      --   type Foo.Bar = ...
-      --   type Baz.Qux = ...
-      --   type Whatever = Whatever Foo.Bar Qux
-      --                            ^^^^^^^ ^^^
-      --                               |    this variable *is* unresolved: it doesn't match any locally-bound type exactly
-      --                               |
-      --                            this variable is *not* unresolved: it matches locally-bound `Foo.Bar` exactly
-      unresolvedVars :: [(v, a)]
-      unresolvedVars =
-        ABT.freeVarOccurrences localVars ty
+bindNames unsafeVarToName nameToVar localVars namespace =
+  -- type is bound here because the where-clause binds a data structure that we only want to compute once, then share
+  -- across all calls to `bindNames` with different types
+  \ty ->
+    let -- Identify the unresolved variables in the type: those whose names aren't an *exact* match for some locally-bound
+        -- type.
+        --
+        -- For example:
+        --
+        --   type Foo.Bar = ...
+        --   type Baz.Qux = ...
+        --   type Whatever = Whatever Foo.Bar Qux
+        --                            ^^^^^^^ ^^^
+        --                               |    this variable *is* unresolved: it doesn't match any locally-bound type exactly
+        --                               |
+        --                            this variable is *not* unresolved: it matches locally-bound `Foo.Bar` exactly
+        unresolvedVars :: [(v, a)]
+        unresolvedVars =
+          ABT.freeVarOccurrences localVars ty
 
-      okTy :: (v, a) -> Names.ResolutionResult a (v, ResolvesTo TypeReference)
-      okTy (v, a) =
-        case Set.size matches of
-          1 -> good (Set.findMin matches)
-          0 -> bad Names.NotFound
-          _ ->
-            let (namespaceMatches, localMatches) =
-                  matches
-                    & Set.toList
-                    & map \case
-                      ResolvesToNamespace ref -> Left ref
-                      ResolvesToLocal name -> Right name
-                    & partitionEithers
-                    & bimap Set.fromList Set.fromList
-             in bad (Names.Ambiguous namespace namespaceMatches localMatches)
-        where
-          matches :: Set (ResolvesTo TypeReference)
-          matches =
-            Names.resolveName (Names.types namespace) (Set.map unsafeVarToName localVars) (unsafeVarToName v)
+        okTy :: (v, a) -> Names.ResolutionResult a (v, ResolvesTo TypeReference)
+        okTy (v, a) =
+          case Set.size matches of
+            1 -> good (Set.findMin matches)
+            0 -> bad Names.NotFound
+            _ ->
+              let (namespaceMatches, localMatches) =
+                    matches
+                      & Set.toList
+                      & map \case
+                        ResolvesToNamespace ref -> Left ref
+                        ResolvesToLocal name -> Right name
+                      & partitionEithers
+                      & bimap Set.fromList Set.fromList
+               in bad (Names.Ambiguous namespace namespaceMatches localMatches)
+          where
+            matches :: Set (ResolvesTo TypeReference)
+            matches =
+              resolveTypeName (unsafeVarToName v)
 
-          bad = Left . Seq.singleton . Names.TypeResolutionFailure (HQ.NameOnly (unsafeVarToName v)) a
-          good = Right . (v,)
-   in List.validate okTy unresolvedVars <&> \resolutions ->
-        let (namespaceResolutions, localResolutions) = partitionResolutions resolutions
-         in ty
-              -- Apply namespace resolutions (replacing "Foo" with #Foo where "Foo" refers to namespace)
-              & bindExternal namespaceResolutions
-              -- Apply local resolutions (replacing "Foo" with "Full.Name.Foo" where "Full.Name.Foo" is in local vars)
-              & ABT.substsInheritAnnotation [(v, Type.var () (nameToVar name)) | (v, name) <- localResolutions]
+            bad = Left . Seq.singleton . Names.TypeResolutionFailure (HQ.NameOnly (unsafeVarToName v)) a
+            good = Right . (v,)
+     in List.validate okTy unresolvedVars <&> \resolutions ->
+          let (namespaceResolutions, localResolutions) = partitionResolutions resolutions
+           in ty
+                -- Apply namespace resolutions (replacing "Foo" with #Foo where "Foo" refers to namespace)
+                & bindExternal namespaceResolutions
+                -- Apply local resolutions (replacing "Foo" with "Full.Name.Foo" where "Full.Name.Foo" is in local vars)
+                & ABT.substsInheritAnnotation [(v, Type.var () (nameToVar name)) | (v, name) <- localResolutions]
+  where
+    resolveTypeName :: Name -> Set (ResolvesTo TypeReference)
+    resolveTypeName =
+      Names.resolveName (Names.types namespace) (Set.map unsafeVarToName localVars)