give the BiMultimap api some love

lib/unison-util-bimultimap/package.yaml

@@ -11,24 +11,37 @@ library:
 dependencies:
   - base
   - containers
+  - nonempty-containers
   - unison-prelude
 
 ghc-options:
   -Wall
 
 default-extensions:
   - ApplicativeDo
+  - BangPatterns
   - BlockArguments
+  - DeriveAnyClass
   - DeriveFunctor
+  - DeriveGeneric
+  - DeriveTraversable
   - DerivingStrategies
+  - DerivingVia
   - DoAndIfThenElse
+  - DuplicateRecordFields
   - FlexibleContexts
   - FlexibleInstances
+  - GeneralizedNewtypeDeriving
   - ImportQualifiedPost
   - LambdaCase
   - MultiParamTypeClasses
   - NamedFieldPuns
+  - OverloadedStrings
+  - PatternSynonyms
+  - RankNTypes
   - ScopedTypeVariables
+  - StandaloneDeriving
   - TupleSections
   - TypeApplications
+  - TypeFamilies
   - ViewPatterns

lib/unison-util-bimultimap/src/Unison/Util/BiMultimap.hs

@@ -1,45 +1,48 @@
-module Unison.Util.BiMultimap where
+-- | A left-unique relation.
+module Unison.Util.BiMultimap
+  ( BiMultimap (..),
+    Unison.Util.BiMultimap.empty,
+
+    -- ** Lookup
+    lookupDom,
+    lookupRan,
+
+    -- ** Sets
+    ran,
+
+    -- ** Insert
+    insert,
+    unsafeInsert,
+  )
+where
 
 import Data.Map.Strict qualified as Map
-import Data.Maybe (fromJust)
 import Data.Set qualified as Set
+import Data.Set.NonEmpty (NESet)
+import Data.Set.NonEmpty qualified as Set.NonEmpty
 import Unison.Prelude
+import Unison.Util.Map qualified as Map
 
-data BiMultimap a b = BiMultimap {toMultimap :: Map a (Set b), toMapR :: Map b a} deriving (Eq, Ord, Show)
+data BiMultimap a b = BiMultimap
+  { toMultimap :: Map a (NESet b),
+    toMapR :: Map b a
+  }
+  deriving (Eq, Ord, Show)
 
+-- | An empty left-unique relation.
 empty :: (Ord a, Ord b) => BiMultimap a b
 empty = BiMultimap mempty mempty
 
--- |
--- >>> insert 1 'a' empty
--- >>> insert 1 'b' (insert 1 'a' empty)
--- >>> insert 2 'b' (insert 1 'a' (insert 1 'b' empty))
--- >>> insert 2 'a' (insert 2 'b' (insert 1 'a' (insert 1 'b' empty)))
--- BiMultimap {toMultimap = fromList [(1,fromList "a")], toMapR = fromList [('a',1)]}
--- BiMultimap {toMultimap = fromList [(1,fromList "ab")], toMapR = fromList [('a',1),('b',1)]}
--- BiMultimap {toMultimap = fromList [(1,fromList "a"),(2,fromList "b")], toMapR = fromList [('a',1),('b',2)]}
--- BiMultimap {toMultimap = fromList [(2,fromList "ab")], toMapR = fromList [('a',2),('b',2)]}
-insert :: (Ord a, Ord b) => a -> b -> BiMultimap a b -> BiMultimap a b
-insert a b m@(BiMultimap l r) =
-  let lInserted = Map.insertWith (<>) a (Set.singleton b) l
-      lDeleted a' = Map.alter (\(fromJust -> s) -> let s' = Set.delete b s in if Set.null s' then Nothing else Just s') a' lInserted
-      rInserted = Map.insert b a r
-   in case Map.lookup b r of
-        Just a' -> if a' == a then m else BiMultimap (lDeleted a') rInserted
-        Nothing -> BiMultimap lInserted rInserted
-
--- | Like @insert x y@, except the caller is responsible for ensuring that @y@ is not already related to a different
--- @x@. If it is, the resulting relation will have an internal structural violation.
-unsafeInsert :: (Ord a, Ord b) => a -> b -> BiMultimap a b -> BiMultimap a b
-unsafeInsert x y (BiMultimap xs ys) =
-  BiMultimap
-    (Map.alter (Just . maybe (Set.singleton y) (Set.insert y)) x xs)
-    (Map.insert y x ys)
-
+-- | Look up the set of @b@ related to an @a@.
+--
+-- /O(log a)/.
 lookupDom :: Ord a => a -> BiMultimap a b -> Set b
 lookupDom a (BiMultimap l _) =
-  fromMaybe Set.empty (Map.lookup a l)
+  maybe Set.empty Set.NonEmpty.toSet (Map.lookup a l)
 
+-- | Look up the @a@ related to a @b@.
+--
+-- /O(log b)/.
 lookupRan :: Ord b => b -> BiMultimap a b -> Maybe a
 lookupRan b (BiMultimap _ r) =
   Map.lookup b r
@@ -50,3 +53,40 @@ lookupRan b (BiMultimap _ r) =
 ran :: BiMultimap a b -> Set b
 ran =
   Map.keysSet . toMapR
+
+-- | Insert a pair into a left-unique relation, maintaining left-uniqueness, preferring the latest inserted element.
+--
+-- That is, if a left-unique relation already contains the pair @(x, y)@, then inserting the pair @(z, y)@ will cause
+-- the @(x, y)@ pair to be deleted.
+insert :: (Ord a, Ord b) => a -> b -> BiMultimap a b -> BiMultimap a b
+insert a b m@(BiMultimap l r) =
+  case Map.alterF (upsertFunc a) b r of
+    (Ignored, _) -> m
+    (Inserted, r') -> BiMultimap l' r'
+    (Replaced old, r') ->
+      let l'' = Map.update (Set.NonEmpty.nonEmptySet . Set.NonEmpty.delete b) old l'
+       in BiMultimap l'' r'
+  where
+    l' = Map.upsert (maybe (Set.NonEmpty.singleton b) (Set.NonEmpty.insert b)) a l
+
+-- @upsertFunc x@ returns a function that upserts @x@, suitable for passing to @Map.alterF@.
+upsertFunc :: Eq a => a -> Maybe a -> (UpsertResult a, Maybe a)
+upsertFunc new existing =
+  case existing of
+    Nothing -> (Inserted, Just new)
+    Just old
+      | old == new -> (Ignored, existing)
+      | otherwise -> (Replaced old, Just new)
+
+data UpsertResult old
+  = Ignored -- Ignored because an equivalent thing was already there
+  | Inserted -- Inserted something new
+  | Replaced old -- Replaced what was there, here's the old thing
+
+-- | Like @insert x y@, except the caller is responsible for ensuring that @y@ is not already related to a different
+-- @x@. If it is, the resulting relation will have an internal structural violation.
+unsafeInsert :: (Ord a, Ord b) => a -> b -> BiMultimap a b -> BiMultimap a b
+unsafeInsert x y (BiMultimap xs ys) =
+  BiMultimap
+    (Map.upsert (maybe (Set.NonEmpty.singleton y) (Set.NonEmpty.insert y)) x xs)
+    (Map.insert y x ys)

lib/unison-util-bimultimap/unison-util-bimultimap.cabal

@@ -22,23 +22,36 @@ library
       src
   default-extensions:
       ApplicativeDo
+      BangPatterns
       BlockArguments
+      DeriveAnyClass
       DeriveFunctor
+      DeriveGeneric
+      DeriveTraversable
       DerivingStrategies
+      DerivingVia
       DoAndIfThenElse
+      DuplicateRecordFields
       FlexibleContexts
       FlexibleInstances
+      GeneralizedNewtypeDeriving
       ImportQualifiedPost
       LambdaCase
       MultiParamTypeClasses
       NamedFieldPuns
+      OverloadedStrings
+      PatternSynonyms
+      RankNTypes
       ScopedTypeVariables
+      StandaloneDeriving
       TupleSections
       TypeApplications
+      TypeFamilies
       ViewPatterns
   ghc-options: -Wall
   build-depends:
       base
     , containers
+    , nonempty-containers
     , unison-prelude
   default-language: Haskell2010

unison-cli/src/Unison/Codebase/Editor/HandleInput/Merge.hs

@@ -16,6 +16,7 @@ import Data.List.NonEmpty (pattern (:|))
 import Data.Map.Strict qualified as Map
 import Data.Semialign (alignWith)
 import Data.Set qualified as Set
+import Data.Set.NonEmpty qualified as Set.NonEmpty
 import Data.Text qualified as Text
 import Data.Text.IO qualified as Text
 import Data.Text.Lazy qualified as Text.Lazy
@@ -702,8 +703,8 @@ relationToLuniqRelation relation =
 --
 -- TODO move this helper to some other module
 luniqRelationToRelation :: forall a b. (Ord a, Ord b) => BiMultimap a b -> Relation a b
-luniqRelationToRelation relation =
-  Relation.fromMultimap (BiMultimap.toMultimap relation)
+luniqRelationToRelation =
+  Relation.fromMultimap . Map.map Set.NonEmpty.toSet . BiMultimap.toMultimap
 
 -- | Return the set of elements that appear in at least two of the given sets.
 duplicates :: forall a. Ord a => [Set a] -> Set a