--
-- Adapted from the program "infer", believed to have been originally
-- authored by Philip Wadler, and used in the nofib benchmark suite
-- since at least the late 90s.
--

module Substitution   (Sub, applySub, lookupSub, emptySub, extendSub,
                       makeSub, thenSub, domSub, unifySub)
                      where

import Type
import FiniteMap
import MaybeM

data  Sub  =  MkSub (FM TVarId MonoType)
rep                           ::  Sub -> FM TVarId MonoType
rep (MkSub f)                 =   f
applySub                      ::  Sub -> MonoType -> MonoType
applySub s (TVar x)           =   lookupSub s x
applySub s (TCon k ts)        =   TCon k (map (applySub s) ts)
lookupSub                     ::  Sub -> TVarId -> MonoType
lookupSub s x                 =   lookupElseFM (TVar x) (rep s) x
unitSub                       ::  TVarId -> MonoType -> Sub
unitSub x t                   =   MkSub (makeFM [(x,t)])
emptySub                      ::  Sub
emptySub                      =   MkSub emptyFM
makeSub                       ::  [(TVarId, MonoType)] -> Sub
makeSub xts                   =   MkSub (makeFM xts)
extendSub                     ::  Sub -> TVarId -> MonoType -> Sub
extendSub s x t               =   s `thenSub` unitSub x (applySub s t)
thenSub                       ::  Sub -> Sub -> Sub
r `thenSub` s                 =   MkSub (mapFM (applySub s) (rep r) `thenFM` rep s)
domSub                        ::  Sub -> [TVarId]
domSub s                      =   domFM (rep s)
unifySub                              =  unify
unify                                 :: MonoType -> MonoType -> Sub -> Maybe Sub
unify (TVar x) u s                    =  unifyTVar x u s
unify t (TVar y) s                    =  unifyTVar y t s
unify (TCon j ts) (TCon k us) s       =  (j == k) `guardM` unifies ts us s
unifies                               :: [MonoType] -> [MonoType] -> Sub -> Maybe Sub
unifies [] [] s                       =  returnM s
unifies (t:ts) (u:us) s               =  unify t u s `thenM` (\s' -> unifies ts us s')
unifyTVar                             :: TVarId -> MonoType -> Sub -> Maybe Sub
unifyTVar x t s | x `elem` domSub s     =  unify (lookupSub s x) t s
                | TVar x == t         =  returnM s
                | x `elem` freeVars t   =  failM
                | otherwise           =  returnM (extendSub s x t)    
freeVars                              =  freeTVarMono