# Unification with sequence variables and flexible arity functions

Today I wrote an implementation of the unification algorithm found in Temur Kutsia's 2002 paper. I didn't just do this for fun, it's related to other research I'm doing. I'm feeling more confident in my Haskell these days, and I was going to go through this code to clean it up myself but, since this is a pretty general and nicely self-contained module, I thought I'd ask for some outside review to see what others would do. I am open to all suggestions, particularly related to the massive transformations function.

Other questions I had: does it look like I can make use of Foldable/Traversable? Does it look worth it to add a Substitutable type class? Am I doing lots of unnecessary work (probably)?

*Note: I changed all variables ending with apostrophes to end with p since it was breaking the syntax highlighter. Sorry if that's ugly.

module Unification where

import Data.Maybe
import Data.List
import qualified Data.Set as Set
import qualified Data.Map as Map
import Debug.Trace

-- a is the set from which both individual and sequence variables are taken
-- c is the set of constant values
data Term a c
= Const c
| IVar a
| SVar a
| Fun a [Term a c]
| Seq [Term a c]
deriving Eq

instance (Show a, Show c) => Show (Term a c) where
show (Const c) = show c
show (IVar a) = show a
show (SVar a) = show a
show (Fun a ts) = show a ++ "(" ++ (concat $intersperse ","$ map show ts) ++ ")"
show (Seq ts) = "[" ++ (concat $intersperse ","$ map show ts) ++ "]"

type Equation a c = (Term a c, Term a c)

data UniTree a c = Success
| Failure
| Problem (Equation a c) [(Substitution a c, UniTree a c)]
deriving (Eq, Show)

type Substitution a c = Map.Map a (Term a c)

vars :: Ord a => Term a c -> Set.Set a
vars (Fun v ts) = Set.insert v $Set.unions (map vars ts) vars (SVar v) = Set.singleton v vars (IVar v) = Set.singleton v vars _ = Set.empty seqVars :: Ord a => Term a c -> Set.Set a seqVars (Fun v ts) = Set.unions (map seqVars ts) seqVars (SVar v) = Set.singleton v seqVars _ = Set.empty powerset :: Ord a => Set.Set a -> [Set.Set a] powerset = map Set.fromAscList . filterM (const [False,True]) . Set.toList substitute :: Ord a => Substitution a c -> Term a c -> Term a c substitute env (IVar v) = Map.findWithDefault (IVar v) v env substitute env (SVar v) = Map.findWithDefault (SVar v) v env substitute env (Fun v ts) = Fun v$ flattenSeq (map (substitute env) ts)
substitute env (Seq ts) =
Seq $flattenSeq (map (substitute env) ts) substitute env t = t substituteSeq :: Ord a => Substitution a c -> [Term a c] -> [Term a c] substituteSeq env ts = flattenSeq$ map (substitute env) ts

flattenSeq :: [Term a c] -> [Term a c]
flattenSeq [] = []
flattenSeq (Seq subs:ts) = subs ++ ts
flattenSeq (t:ts) = t:flattenSeq ts

composeUnifySubs :: Ord a => Substitution a c -> Substitution a c -> Substitution a c
composeUnifySubs e1 e2 =
let e1p = Map.filterWithKey cleanSubs $Map.map (substitute e2) e1 -- Relies on left-biased implementation of Map.union to remove overlapped elements from e2 in e1p Map.union e2 where cleanSubs v1 (IVar v2) = not$ v1 == v2
cleanSubs v1 (Seq [SVar v2]) = not $v1 == v2 cleanSubs _ _ = True unifyTree :: (Show a, Ord a, Show c, Eq c) => State [a] a -> Term a c -> Term a c -> UniTree a c unifyTree getFresh t1 t2 = Problem (t1p,t2p) [(env, transformed env) | (env,_) <- ts] where (Problem (t1p,t2p) ts) = projections t1 t2 transformed env = evalState (transformations getFresh (substitute env t1p) (substitute env t2p)) (Set.toList$ vars t1 Set.union vars t2)

-- only used for the first level of the tree
projections :: (Ord a, Eq c) => Term a c -> Term a c -> UniTree a c
projections t1 t2 =
Problem (t1,t2) [(env, Success) | env <- emptyEnvs]
where svars = seqVars t1 Set.union seqVars t2
emptyEnvs = map toEmpties $powerset svars toEmpties = Map.fromSet (const (Seq [])) -- used for all levels after the first level of the tree transformations :: (Show a, Ord a, Show c, Eq c) => State [a] a -> Term a c -> Term a c -> State [a] (UniTree a c) transformations _ t1@(IVar _) t2@(IVar _) | t1 == t2 = return$ Problem (t1,t2) [(Map.empty, Success)]
transformations _ t1@(Const _) t2@(Const _)
| t1 == t2 = return $Problem (t1,t2) [(Map.empty, Success)] | otherwise = return Failure transformations _ t1@(Fun _ _) t2@(Fun _ _) | t1 == t2 = return$ Problem (t1,t2) [(Map.empty, Success)]
transformations _ (Fun _ _) (Const _) = return Failure
transformations _ (Const _) (Fun _ _) = return Failure
transformations _ t1@(IVar v) t2
| v Set.member vars t2 = return Failure
| otherwise = return $Problem (t1,t2) [(Map.singleton v t2, Success)] transformations _ t1 t2@(IVar v) | v Set.member vars t1 = return Failure | otherwise = return$ Problem (t1,t2) [(Map.singleton v t1, Success)]
transformations _ (Fun v1 _) (Fun v2 _)
| v1 /= v2 = return Failure
transformations _ (Fun _ []) (Fun _ (t:ts)) = return Failure
transformations _ (Fun _ (t:ts)) (Fun _ []) = return Failure
transformations _ (Fun _ (t1@(SVar v):t1s)) (Fun _ (t2:t2s))
| t1 /= t2 && v Set.member seqVars t2 = return Failure
transformations _ (Fun _ (t1:t1s)) (Fun _ (t2@(SVar v):t2s))
| t1 /= t2 && v Set.member seqVars t1 = return Failure
transformations getFresh (Fun f (t1@(SVar v1):t1s)) (Fun _ (t2@(SVar v2):t2s))
| v1 == v2 = do
transformed <- transformations getFresh (Fun f t1s) (Fun f t2s)
return $Problem (t1,t2) [(Map.empty, transformed)] | otherwise = do v1p <- getFresh v2p <- getFresh let sub1 = Map.singleton v1 (Seq [t2]) sub2 = Map.singleton v1 (Seq [t2, (SVar v1p)]) sub3 = Map.singleton v2 (Seq [t1, (SVar v2p)]) trans1 <- transformations getFresh (Fun f (substituteSeq sub1 t1s)) (Fun f (substituteSeq sub1 t2s)) trans2 <- transformations getFresh (Fun f (SVar v1p:(substituteSeq sub2 t1s))) (Fun f (substituteSeq sub2 t2s)) trans3 <- transformations getFresh (Fun f (substituteSeq sub3 t1s)) (Fun f (SVar v2p:(substituteSeq sub3 t2s))) return$ Problem (t1,t2) [(sub1, trans1),(sub2, trans2),(sub3, trans3)]
transformations getFresh (Fun f (t1@(SVar v):t1s)) (Fun _ (t2:t2s))
| not $v Set.member vars t2 = do vp <- getFresh let sub1 = Map.singleton v (Seq [t2]) sub2 = Map.singleton v (Seq [t2,(SVar vp)]) trans1 <- transformations getFresh (Fun f (substituteSeq sub1 t1s)) (Fun f (substituteSeq sub1 t2s)) trans2 <- transformations getFresh (Fun f (SVar vp:(substituteSeq sub2 t1s))) (Fun f (substituteSeq sub2 t2s)) return$ Problem (t1,t2) [(sub1, trans1),(sub2, trans2)]
transformations getFresh (Fun f (t1:t1s)) (Fun _ (t2@(SVar v):t2s))
| not $v Set.member vars t1 = do vp <- getFresh let sub1 = Map.singleton v (Seq [t1]) sub2 = Map.singleton v (Seq [t1,(SVar vp)]) trans1 <- transformations getFresh (Fun f (substituteSeq sub1 t1s)) (Fun f (substituteSeq sub1 t2s)) trans2 <- transformations getFresh (Fun f (substituteSeq sub2 t1s)) (Fun f (SVar vp:(substituteSeq sub2 t2s))) return$ Problem (t1,t2) [(sub1, trans1), (sub2, trans2)]
transformations getFresh (Fun f (t1:t1s)) (Fun _ (t2:t2s)) = do
transformed <- transformations getFresh t1 t2
case transformed of
Failure -> return Failure
Problem _ [(sub, Success)] -> do
transformed <- transformations getFresh (Fun f (substituteSeq sub t1s)) (Fun f (substituteSeq sub t2s))
return $Problem (t1,t2) [(sub, transformed)] Problem _ rs -> do alltransformed <- mapM doSub rs return$ Problem (t1,t2) alltransformed
where doSub (sub, Problem (q,r) _) = do
let t1sp = substituteSeq sub t1s
t2sp = substituteSeq sub t2s
transformed <- transformations getFresh (Fun f (q:t1sp)) (Fun f (r:t2sp))
return (sub, transformed)
doSub (sub, Failure) = return (sub, Failure)
doSub (sub, Success) = return (sub, Success)
transformations _ t1 t2 = error $"got: " ++ show t1 ++ " " ++ show t2 unifyMapsFromTree :: Ord a => UniTree a c -> [Substitution a c] unifyMapsFromTree = catMaybes . unifyMaybeMaps Map.empty unifyMaybeMaps :: Ord a => Substitution a c -> UniTree a c -> [Maybe (Substitution a c)] unifyMaybeMaps _ Failure = [Nothing] unifyMaybeMaps env Success = [Just env] unifyMaybeMaps env (Problem _ ts) = concat$ map (unifyMaybeMap env) ts

unifyMaybeMap :: Ord a => Substitution a c -> (Substitution a c, UniTree a c) -> [Maybe (Substitution a c)]
unifyMaybeMap env (_, Failure) = [Nothing]
unifyMaybeMap env (envp, Success) = [Just $composeUnifySubs env envp] unifyMaybeMap env (envp, Problem _ ts) = concat$ map (unifyMaybeMap (composeUnifySubs env envp)) ts

unifiers :: (Show a, Ord a, Show c, Eq c) => State [a] a -> Term a c -> Term a c -> [Substitution a c]
unifiers getFresh t1 t2 =
let tree = unifyTree getFresh t1 t2
vs = vars t1 Set.union vars t2
in map (flattenNew vs) $unifyMapsFromTree tree flattenNew :: (Show a, Ord a, Show c) => Set.Set a -> Substitution a c -> Substitution a c flattenNew names sub = Map.difference substNew new where new = Map.difference sub$ Map.fromSet (const (Seq [])) names
substNew = Map.map (substitute new) sub