Rewrite patchThatChangesMap

The new implementation doesn't use `undefined`, and hopefully works better

Author: ryantrinkle

src/Data/Patch/MapWithMove.hs

@@ -16,9 +16,8 @@ module Data.Patch.MapWithMove where
 
 import Data.Patch.Class
 
-import Control.Arrow
 import Control.Lens hiding (from, to)
-import Control.Monad.Trans.State
+import Data.Align
 import Data.Foldable
 import Data.Function
 import Data.List
@@ -28,9 +27,9 @@ import Data.Maybe
 #if !MIN_VERSION_base(4,11,0)
 import Data.Semigroup (Semigroup (..))
 #endif
+import Data.Set (Set)
 import qualified Data.Set as Set
 import Data.These (These(..))
-import Data.Tuple
 
 -- | Patch a Map with additions, deletions, and moves.  Invariant: If key @k1@
 -- is coming from @From_Move k2@, then key @k2@ should be going to @Just k1@,
@@ -185,37 +184,36 @@ patchThatChangesAndSortsMapWith cmp oldByIndex newByIndexUnsorted = patchThatCha
 
 -- | Create a 'PatchMapWithMove' that, if applied to the first 'Map' provided,
 -- will produce the second 'Map'.
-patchThatChangesMap :: (Ord k, Ord v) => Map k v -> Map k v -> PatchMapWithMove k v
+patchThatChangesMap :: forall k v. (Ord k, Ord v) => Map k v -> Map k v -> PatchMapWithMove k v
 patchThatChangesMap oldByIndex newByIndex = patch
-  where oldByValue = Map.fromListWith Set.union $ swap . first Set.singleton <$> Map.toList oldByIndex
-        (insertsAndMoves, unusedValuesByValue) = flip runState oldByValue $ do
-          let f k v = do
-                remainingValues <- get
-                let putRemainingKeys remainingKeys = put $ if Set.null remainingKeys
-                      then Map.delete v remainingValues
-                      else Map.insert v remainingKeys remainingValues
-                case Map.lookup v remainingValues of
-                  Nothing -> return $ NodeInfo (From_Insert v) $ Just undefined -- There's no existing value we can take
-                  Just fromKs ->
-                    if k `Set.member` fromKs
-                    then do
-                      putRemainingKeys $ Set.delete k fromKs
-                      return $ NodeInfo (From_Move k) $ Just undefined -- There's an existing value, and it's here, so no patch necessary
-                    else do
-                      (fromK, remainingKeys) <- return $
-                        fromMaybe (error "PatchMapWithMove.patchThatChangesMap: impossible: fromKs was empty") $
-                        Set.minView fromKs -- There's an existing value, but it's not here; move it here
-                      putRemainingKeys remainingKeys
-                      return $ NodeInfo (From_Move fromK) $ Just undefined
-          Map.traverseWithKey f newByIndex
-        unusedOldKeys = fold unusedValuesByValue
-        pointlessMove k = \case
-          From_Move k' | k == k' -> True
-          _ -> False
-        keyWasMoved k = if k `Map.member` oldByIndex && not (k `Set.member` unusedOldKeys)
-          then Just undefined
-          else Nothing
-        patch = unsafePatchMapWithMove $ Map.filterWithKey (\k -> not . pointlessMove k . _nodeInfo_from) $ Map.mergeWithKey (\k a _ -> Just $ nodeInfoSetTo (keyWasMoved k) a) (Map.mapWithKey $ \k -> nodeInfoSetTo $ keyWasMoved k) (Map.mapWithKey $ \k _ -> NodeInfo From_Delete $ keyWasMoved k) insertsAndMoves oldByIndex
+  where invert :: Map k v -> Map v (Set k)
+        invert = Map.fromListWith (<>) . fmap (\(k, v) -> (v, Set.singleton k)) . Map.toList
+        -- In the places where we use unionDistinct, a non-distinct key indicates a bug in this function
+        unionDistinct :: forall k' v'. Ord k' => Map k' v' -> Map k' v' -> Map k' v'
+        unionDistinct = Map.unionWith (error "patchThatChangesMap: non-distinct keys")
+        unionPairDistinct :: (Map k (From k v), Map k (To k)) -> (Map k (From k v), Map k (To k)) -> (Map k (From k v), Map k (To k))
+        unionPairDistinct (oldFroms, oldTos) (newFroms, newTos) = (unionDistinct oldFroms newFroms, unionDistinct oldTos newTos)
+        -- Generate patch info for a single value
+        -- Keys that are found in both the old and new sets will not be patched
+        -- Keys that are found in only the old set will be moved to a new position if any are available; otherwise they will be deleted
+        -- Keys that are found in only the new set will be populated by moving an old key if any are available; otherwise they will be inserted
+        patchSingleValue :: v -> Set k -> Set k -> (Map k (From k v), Map k (To k))
+        patchSingleValue v oldKeys newKeys = foldl' unionPairDistinct mempty $ align (toList $ oldKeys `Set.difference` newKeys) (toList $ newKeys `Set.difference` oldKeys) <&> \case
+          This oldK -> (mempty, Map.singleton oldK Nothing) -- There's nowhere for this value to go, so we know we are deleting it
+          That newK -> (Map.singleton newK $ From_Insert v, mempty) -- There's nowhere fo this value to come from, so we know we are inserting it
+          These oldK newK -> (Map.singleton newK $ From_Move oldK, Map.singleton oldK $ Just newK)
+        -- Run patchSingleValue on a These.  Missing old or new sets are considered empty
+        patchSingleValueThese :: v -> These (Set k) (Set k) -> (Map k (From k v), Map k (To k))
+        patchSingleValueThese v = \case
+          This oldKeys -> patchSingleValue v oldKeys mempty
+          That newKeys -> patchSingleValue v mempty newKeys
+          These oldKeys newKeys -> patchSingleValue v oldKeys newKeys
+        -- Generate froms and tos for all values, then merge them together
+        (froms, tos) = foldl' unionPairDistinct mempty $ Map.mapWithKey patchSingleValueThese $ align (invert oldByIndex) (invert newByIndex)
+        patch = unsafePatchMapWithMove $ align froms tos <&> \case
+          This from -> NodeInfo from Nothing -- Since we don't have a 'to' record for this key, that must mean it isn't being moved anywhere, so it should be deleted.
+          That to -> NodeInfo From_Delete to -- Since we don't have a 'from' record for this key, it must be getting deleted
+          These from to -> NodeInfo from to
 
 -- | Change the 'From' value of a 'NodeInfo'
 nodeInfoMapFrom :: (From k v -> From k v) -> NodeInfo k v -> NodeInfo k v