make Nametree its own data type

unisonweb · Oct 25, 2023 · 4943801 · 4943801
1 parent 9e21bfa
commit 4943801
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 44 deletions.
diff --git a/lib/unison-util-nametree/package.yaml b/lib/unison-util-nametree/package.yaml
@@ -7,7 +7,7 @@ ghc-options: -Wall
 dependencies:
   - base
   - containers
-  - free
+  - lens
   - semigroups
   - unison-core
   - unison-core1

diff --git a/lib/unison-util-nametree/src/Unison/Util/Nametree.hs b/lib/unison-util-nametree/src/Unison/Util/Nametree.hs
@@ -1,6 +1,6 @@
 module Unison.Util.Nametree
   ( -- * Nametree
-    Nametree,
+    Nametree (..),
     traverseNametreeWithName,
     flattenNametree,
     unflattenNametree,
@@ -12,7 +12,7 @@ module Unison.Util.Nametree
   )
 where
 
-import Control.Comonad.Cofree (Cofree ((:<)))
+import Control.Lens ((^.))
 import Data.List.NonEmpty (NonEmpty, pattern (:|))
 import Data.Map.Merge.Strict qualified as Map
 import Data.Map.Strict qualified as Map
@@ -24,47 +24,50 @@ import Unison.Prelude
 import Unison.Util.BiMultimap (BiMultimap)
 import Unison.Util.BiMultimap qualified as BiMultimap
 
--- | A nametree has values, and a collection of children nametrees keyed by name segment.
-type Nametree a =
-  Cofree (Map NameSegment) a
+-- | A nametree has a value, and a collection of children nametrees keyed by name segment.
+data Nametree a = Nametree
+  { value :: !a,
+    children :: !(Map NameSegment (Nametree a))
+  }
+  deriving stock (Functor, Generic, Show)
 
 -- | Traverse over a nametree, with access to the list of name segments (in reverse order) leading to each value.
 traverseNametreeWithName :: Applicative f => ([NameSegment] -> a -> f b) -> Nametree a -> f (Nametree b)
 traverseNametreeWithName f =
   go []
   where
-    go names (x :< xs) =
-      (:<) <$> f names x <*> Map.traverseWithKey (\name -> go (name : names)) xs
+    go names (Nametree x xs) =
+      Nametree <$> f names x <*> Map.traverseWithKey (\name -> go (name : names)) xs
 
-mergeNametrees ::
-  (a -> c) ->
-  (b -> c) ->
-  (a -> b -> c) ->
-  Nametree a ->
-  Nametree b ->
-  Nametree c
+mergeNametrees :: (a -> c) -> (b -> c) -> (a -> b -> c) -> Nametree a -> Nametree b -> Nametree c
 mergeNametrees ac bc abc =
-  let go (a :< as) (b :< bs) =
-        abc a b
-          :< Map.merge
-            (Map.mapMaybeMissing (\_nameSeg cofreeA -> Just (ac <$> cofreeA)))
-            (Map.mapMaybeMissing (\_nameSeg cofreeB -> Just (bc <$> cofreeB)))
-            (Map.zipWithMaybeMatched (\_nameSeg cofreeA cofreeB -> Just (go cofreeA cofreeB)))
+  go
+  where
+    go (Nametree a as) (Nametree b bs) =
+      Nametree
+        (abc a b)
+        ( Map.merge
+            (Map.mapMissing (\_ -> fmap ac))
+            (Map.mapMissing (\_ -> fmap bc))
+            (Map.zipWithMatched (\_ -> go))
             as
             bs
-   in go
+        )
 
 zipNametrees :: (a -> b -> c) -> Nametree a -> Nametree b -> Nametree c
 zipNametrees f =
-  let go (a :< as) (b :< bs) =
-        f a b
-          :< Map.merge
+  go
+  where
+    go (Nametree a as) (Nametree b bs) =
+      Nametree
+        (f a b)
+        ( Map.merge
             Map.dropMissing
             Map.dropMissing
-            (Map.zipWithMaybeMatched (\_nameSeg cofreeA cofreeB -> Just (go cofreeA cofreeB)))
+            (Map.zipWithMatched (\_ -> go))
             as
             bs
-   in go
+        )
 
 -- | 'flattenNametree' organizes a nametree like
 --
@@ -95,8 +98,8 @@ flattenNametree =
       [NameSegment] ->
       Nametree (Defns (Map NameSegment terms) (Map NameSegment types)) ->
       Defns (BiMultimap terms Name) (BiMultimap types Name)
-    go prefix (defns :< children) =
-      foldr step (fff defns) (Map.toList children)
+    go prefix (Nametree defns0 children) =
+      foldr step defns1 (Map.toList children)
       where
         step ::
           (NameSegment, Nametree (Defns (Map NameSegment terms) (Map NameSegment types))) ->
@@ -107,18 +110,16 @@ flattenNametree =
            in -- These unions are safe because the input nametree had unconflicted names
               Defns (BiMultimap.unsafeUnion accTerms childTerms) (BiMultimap.unsafeUnion accTypes childTypes)
 
-        fff ::
-          Defns (Map NameSegment terms) (Map NameSegment types) ->
-          Defns (BiMultimap terms Name) (BiMultimap types Name)
-        fff Defns {terms, types} =
+        defns1 :: Defns (BiMultimap terms Name) (BiMultimap types Name)
+        defns1 =
           Defns
             { terms =
                 BiMultimap.fromRange
                   ( Map.mapKeysMonotonic
                       (\name -> Name.fromReverseSegments (name :| prefix))
-                      terms
+                      (defns0 ^. #terms)
                   ),
-              types = BiMultimap.fromRange (Map.mapKeysMonotonic Name.fromSegment types)
+              types = BiMultimap.fromRange (Map.mapKeysMonotonic Name.fromSegment (defns0 ^. #types))
             }
 
 unflattenNametree ::
@@ -158,7 +159,7 @@ unflattenNametree defns0 =
       unflatten a b =
         let (curr, children) = unflattenLevel a b
             finalChildren = fmap (uncurry unflatten) children
-         in curr :< finalChildren
+         in Nametree curr finalChildren
    in unflatten inputTerms inputTypes
 
 -- | Definitions (terms and types) in a namespace.

diff --git a/lib/unison-util-nametree/unison-util-nametree.cabal b/lib/unison-util-nametree/unison-util-nametree.cabal
@@ -56,7 +56,7 @@ library
   build-depends:
       base
     , containers
-    , free
+    , lens
     , semigroups
     , unison-core
     , unison-core1

diff --git a/unison-cli/src/Unison/Codebase/Editor/HandleInput/Merge2.hs b/unison-cli/src/Unison/Codebase/Editor/HandleInput/Merge2.hs
@@ -118,7 +118,7 @@ import Unison.Util.Cache qualified as Cache
 import Unison.Util.Map qualified as Map
 import Unison.Util.Nametree
   ( Defns (..),
-    Nametree,
+    Nametree (..),
     flattenNametree,
     mergeNametrees,
     traverseNametreeWithName,
@@ -601,7 +601,7 @@ namespaceToV3Branch ::
   MergeDatabase ->
   Nametree (Defns (Map NameSegment Referent) (Map NameSegment TypeReference), [CausalHash]) ->
   BranchV3 Transaction
-namespaceToV3Branch db ((Defns {terms, types}, _causalParents) :< children) =
+namespaceToV3Branch db (Nametree (Defns {terms, types}, _causalParents) children) =
   BranchV3.BranchV3
     { terms,
       types,
@@ -612,11 +612,11 @@ namespaceToV3Causal ::
   MergeDatabase ->
   Nametree (Defns (Map NameSegment Referent) (Map NameSegment TypeReference), [CausalHash]) ->
   BranchV3.CausalBranchV3 Transaction
-namespaceToV3Causal db@MergeDatabase {loadCausal} namespace@((_, causalParentHashes) :< _) =
+namespaceToV3Causal db@MergeDatabase {loadCausal} namespace =
   HashHandle.mkCausal
     v2HashHandle
     (HashHandle.hashBranchV3 v2HashHandle v3Branch)
-    (Map.fromList (map (\ch -> (ch, loadCausal ch)) causalParentHashes))
+    (Map.fromList (map (\ch -> (ch, loadCausal ch)) (snd (namespace ^. #value))))
     (pure v3Branch)
   where
     v3Branch :: BranchV3 Transaction
@@ -852,21 +852,23 @@ loadNamespaceInfo0 :: Monad m => Branch m -> CausalHash -> m NamespaceInfo0
 loadNamespaceInfo0 branch causalHash = do
   let terms = Map.map Map.keysSet (branch ^. #terms)
   let types = Map.map Map.keysSet (branch ^. #types)
+  let value = (Defns {terms, types}, causalHash)
   children <-
     for (Map.delete Name.libSegment (branch ^. #children)) \childCausal -> do
       childBranch <- Causal.value childCausal
       loadNamespaceInfo0_ childBranch (childCausal ^. #causalHash)
-  pure ((Defns {terms, types}, causalHash) :< children)
+  pure Nametree {value, children}
 
 loadNamespaceInfo0_ :: Monad m => Branch m -> CausalHash -> m NamespaceInfo0
 loadNamespaceInfo0_ branch causalHash = do
   let terms = Map.map Map.keysSet (branch ^. #terms)
   let types = Map.map Map.keysSet (branch ^. #types)
+  let value = (Defns {terms, types}, causalHash)
   children <-
     for (branch ^. #children) \childCausal -> do
       childBranch <- Causal.value childCausal
       loadNamespaceInfo0_ childBranch (childCausal ^. #causalHash)
-  pure ((Defns {terms, types}, causalHash) :< children)
+  pure Nametree {value, children}
 
 type NamespaceInfo1 =
   Nametree
@@ -989,7 +991,7 @@ checkDeclCoherency loadNumConstructors =
       [NameSegment] ->
       NamespaceInfo1 ->
       StateT DeclCoherencyCheckState (ExceptT Merge.PreconditionViolation m) ()
-    go prefix ((Defns {terms, types}, _) :< children) = do
+    go prefix (Nametree (Defns {terms, types}, _) children) = do
       for_ (Map.toList terms) \case
         (_, Referent.Ref _) -> pure ()
         (_, Referent.Con (ReferenceBuiltin _) _) -> pure ()