Add initial version of unification monad transformer

Yes, it's in the same file as the syntax. Make it work,
make it right, make it fast.

spago.dhall

@@ -11,7 +11,20 @@ When creating a new Spago project, you can use
 to generate this file without the comments in this block.
 -}
 { name = "bergamot"
-, dependencies = [ "foldable-traversable", "lists", "logict", "prelude" ]
+, dependencies =
+  [ "bifunctors"
+  , "control"
+  , "foldable-traversable"
+  , "lazy"
+  , "lists"
+  , "logict"
+  , "maybe"
+  , "newtype"
+  , "ordered-collections"
+  , "prelude"
+  , "transformers"
+  , "tuples"
+  ]
 , packages = ./packages.dhall
 , sources = [ "src/**/*.purs", "test/**/*.purs" ]
 }

src/Language/Bergamot/Syntax.purs

@@ -1,11 +1,34 @@
 module Language.Bergamot.Syntax where
 
-import Data.List (List)
-import Data.Functor (class Functor, (<$>))
+import Prelude (Unit, ($), (<<<), unit, (/=), const, flip, (+))
+
+import Control.Plus (class Plus, empty)
+import Control.Monad (class Monad)
+import Control.Monad.State.Trans (StateT, runStateT)
+import Control.Monad.State.Class (gets, modify)
+import Control.Apply (class Apply, apply)
+import Control.Alt (class Alt, alt)
+import Control.Alternative (class Alternative)
+import Control.Applicative (class Applicative, pure)
+import Control.Bind (class Bind, bind, (>>=))
+import Control.MonadPlus (class MonadPlus)
+import Control.Monad.Logic.Class (class MonadLogic, msplit, interleave)
+import Data.List (List(..), (:))
+import Data.Functor (class Functor, (<$>), map)
 import Data.Eq (class Eq)
-import Data.Traversable (class Traversable, sequence)
-import Data.Foldable (class Foldable, foldr, foldl, foldMap)
+import Data.Ord (class Ord)
+import Data.Traversable (class Traversable, sequence, traverse)
+import Data.Foldable (class Foldable, foldr, foldl, foldMap, any)
 import Data.Monoid ((<>), mempty)
+import Data.Map (Map, lookup, insert)
+import Data.Map as Map
+import Data.Set (Set, singleton, union)
+import Data.Tuple (Tuple(..), fst)
+import Data.Tuple.Nested ((/\))
+import Data.Lazy (Lazy, defer, force)
+import Data.Newtype (class Newtype, un, over2)
+import Data.Maybe (Maybe(..), fromMaybe, isJust)
+import Data.Bifunctor (rmap)
 
 data Expr v
     = Var v
@@ -29,3 +52,136 @@ instance Traversable Expr where
     sequence (Atom name fs) = Atom name <$> sequence (sequence <$> fs)
 
     traverse f e = sequence (f <$> e)
+
+data Stream k = StreamCons k (Lazy (Stream k))
+
+pop :: forall k. Stream k -> Tuple k (Stream k)
+pop (StreamCons k lks) = Tuple k (force lks)
+
+class Ord k <= UnificationVariable k where
+    variables :: Stream k
+
+newtype IntVar = MkIntVar Int
+
+derive instance Eq IntVar
+derive instance Ord IntVar
+
+instance UnificationVariable IntVar where
+    variables = mkVarList 0
+        where mkVarList n = StreamCons (MkIntVar n) $ defer $ \_ -> mkVarList (n+1)
+
+data ComparisonAction k f
+    = Merge k k
+    | Store k (f k)
+    | Fail
+
+class (UnificationVariable k, Traversable f) <= Unifiable k f where
+    variable :: k -> f k
+    squash :: f (f k) -> f k
+    alongside :: f k -> f k -> f (ComparisonAction k f)
+
+instance UnificationVariable k => Unifiable k Expr where
+    variable = Var
+    squash (Var f) = f
+    squash (Atom name args) = Atom name $ squash <$> args
+    alongside (Var k1) (Var k2) = Var (Merge k1 k2)
+    alongside (Var k) f = Var (Store k f)
+    alongside f (Var k) = Var (Store k f)
+    alongside (Atom n1 _) (Atom n2 _) | n1 /= n2 = Var Fail
+    alongside (Atom n1 args1) (Atom _ args2) = Atom n1 $ combine args1 args2
+        where
+          combine Nil Nil = Nil
+          combine (_:_) Nil = (Var Fail) : Nil
+          combine Nil (_:_) = (Var Fail) : Nil
+          combine (x:xs) (y:ys) = alongside x y : combine xs ys
+
+class (Unifiable k f, MonadPlus m) <= MonadUnify k f m | m -> k, m -> f where
+    fresh :: m k
+    merge :: k -> k -> m Unit
+    store :: k -> f k -> m Unit
+    reify :: f k -> m (f k)
+
+unify :: forall k f m. MonadUnify k f m => f k -> f k -> m Unit
+unify f1 f2 =
+    do
+       _ <- traverse process $ alongside f1 f2
+       pure unit
+    where
+      process (Merge k1 k2) = merge k1 k2
+      process (Store k f) = store k f
+      process Fail = empty
+
+type UnificationState k f =
+    { boundVariables :: Map k { equivalence :: Set k, boundTo :: Maybe (f k) }
+    , currentVariables :: Stream k
+    }
+
+newtype UnifyT k f m a = MkUnifyT (StateT (UnificationState k f) m a)
+
+derive instance Newtype (UnifyT k f m a) _
+derive instance Functor m => Functor (UnifyT k f m)
+
+instance Monad m => Apply (UnifyT k f m) where
+    apply m1 m2 = MkUnifyT $ apply (un MkUnifyT m1) (un MkUnifyT m2)
+
+instance Monad m => Applicative (UnifyT k f m) where
+    pure a = MkUnifyT $ pure a
+
+instance Monad m => Bind (UnifyT k f m) where
+    bind m f = MkUnifyT $ (un MkUnifyT m) >>= (un MkUnifyT <<< f)
+
+instance Monad m => Monad (UnifyT k f m)
+
+instance (Monad m, Alt m) => Alt (UnifyT k f m) where
+    alt = over2 MkUnifyT alt
+
+instance (Monad m, Plus m) => Plus (UnifyT k f m) where
+    empty = MkUnifyT empty
+
+instance (Monad m, Alternative m) => Alternative (UnifyT k f m)
+instance (MonadPlus m) => MonadPlus (UnifyT k f m)
+
+instance MonadLogic m => MonadLogic (UnifyT k f m) where
+    msplit m = MkUnifyT $ map (map (rmap MkUnifyT)) $ msplit $ un MkUnifyT m
+    interleave m1 m2 = over2 MkUnifyT interleave m1 m2
+
+instance (Unifiable k f, MonadPlus m) => MonadUnify k f (UnifyT k f m) where
+    fresh = MkUnifyT $ do
+        Tuple k restVariables <- gets (pop <<< _.currentVariables)
+        _ <- modify $ _ { currentVariables = restVariables }
+        pure k
+    merge k1 k2 = do
+        boundVariables <- MkUnifyT $ gets _.boundVariables
+        let
+            equivalence k = fromMaybe (singleton k) (_.equivalence <$> lookup k boundVariables)
+            boundTo k = lookup k boundVariables >>= _.boundTo
+            fullSet = equivalence k1 `union` equivalence k2
+        newTerm <-
+            case boundTo k1 /\ boundTo k2 of
+                Just t1 /\ Just t2 -> unify t1 t2 >>= const (Just <$> reify t1)
+                Just t1 /\ Nothing -> pure $ Just t1
+                Nothing /\ Just t2 -> pure $ Just t2
+                Nothing /\ Nothing -> pure $ Nothing
+        let newMapValue = {equivalence: fullSet, boundTo: newTerm}
+        _ <- MkUnifyT $ modify $ _ { boundVariables = foldr (flip insert newMapValue) boundVariables fullSet }
+        pure unit
+    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 newMapValue = {equivalence: fullSet, boundTo: Just f}
+                _ <- MkUnifyT $ modify $ _ { boundVariables = foldr (flip insert newMapValue) boundVariables fullSet }
+                pure unit
+    reify f =
+        MkUnifyT $ do
+            boundVariables <- gets _.boundVariables
+            let
+                reify' f' = squash $ process <$> f'
+                process k = fromMaybe (variable k) (reify' <$> (lookup k boundVariables >>= _.boundTo))
+            pure $ reify' f
+
+runUnifyT :: forall k f m a. Monad m => UnificationVariable k => UnifyT k f m a -> m a
+runUnifyT m = fst <$> runStateT (un MkUnifyT m) { boundVariables: Map.empty, currentVariables: variables }