Allow store to work if a term is already stored

src/Control/Monad/Unify/Trans.purs

@@ -2,7 +2,7 @@ module Control.Monad.Unify.Trans where
 
 import Prelude (($), (<<<), const, flip, unit)
 
-import Control.Plus (class Plus, empty)
+import Control.Plus (class Plus)
 import Control.Monad (class Monad)
 import Control.Monad.State.Trans (StateT, runStateT)
 import Control.Monad.State.Class (gets, modify)
@@ -14,14 +14,15 @@ import Control.Applicative (class Applicative, pure)
 import Control.Bind (class Bind, bind, (>>=))
 import Control.MonadPlus (class MonadPlus)
 import Data.Functor (class Functor, (<$>))
-import Data.Foldable (any, foldr)
+import Data.Foldable (foldMap, foldr)
 import Data.Map (Map, lookup, insert)
 import Data.Map as Map
 import Data.Set (Set, singleton, union)
 import Data.Tuple (Tuple(..))
 import Data.Tuple.Nested ((/\))
 import Data.Newtype (class Newtype, un)
-import Data.Maybe (Maybe(..), fromMaybe, isJust)
+import Data.Maybe (Maybe(..), fromMaybe)
+import Data.Maybe.First (First(..))
 
 type UnificationState k f =
     { boundVariables :: Map k { equivalence :: Set k, boundTo :: Maybe (f k) }
@@ -64,10 +65,10 @@ instance (Unifiable k f, MonadPlus m) => MonadUnify k f (UnifyT k f m) where
     store k f = do
         boundVariables <- MkUnifyT $ gets _.boundVariables
         let fullSet = fromMaybe (singleton k) (_.equivalence <$> lookup k boundVariables)
-        let anyBound = any (isJust <<< (_>>=(_.boundTo)) <<< (flip lookup boundVariables)) fullSet
-        if anyBound
-            then empty
-            else do
+        let firstBound = un First $ foldMap (First <<< (_>>=(_.boundTo)) <<< (flip lookup boundVariables)) fullSet
+        case firstBound of
+            Just f' -> unify f f'
+            Nothing -> do
                 let newMapValue = {equivalence: fullSet, boundTo: Just f}
                 _ <- MkUnifyT $ modify $ _ { boundVariables = foldr (flip insert newMapValue) boundVariables fullSet }
                 pure unit