{-# LANGUAGE MultiWayIf #-}
module MatchSigs.Matching.Env
  ( Env
  , (/\)
  , checkOr
  , checkAnd
  , introVars
  , tryAssignVar
  , initEnv
  ) where

import           Control.Monad.State.Strict
import           Data.List
import qualified Data.IntMap.Strict as IM

import           MatchSigs.Sig (FreeVarIdx)

type Level = Int
type VarLevel = IM.IntMap Int
type VarAssign = IM.IntMap FreeVarIdx

-- | Context for matching the free vars in two 'Sigs'
data Env =
  MkEnv { Env -> Level
level    :: !Level -- current var level
        , Env -> VarAssign
vass     :: !VarAssign -- map from B vars to A vars
        , Env -> VarAssign
vlA      :: !VarLevel -- the level at which an A var with introduced
        , Env -> VarAssign
vlB      :: !VarLevel
        }

initEnv :: Env
initEnv :: Env
initEnv =
  MkEnv :: Level -> VarAssign -> VarAssign -> VarAssign -> Env
MkEnv { level :: Level
level    = Level
0
        , vass :: VarAssign
vass     = VarAssign
forall a. Monoid a => a
mempty
        , vlA :: VarAssign
vlA      = VarAssign
forall a. Monoid a => a
mempty
        , vlB :: VarAssign
vlB      = VarAssign
forall a. Monoid a => a
mempty
        }

-- | Identify var from one sig with var in other sig
tryAssignVar :: FreeVarIdx
             -> FreeVarIdx
             -> State Env Bool
tryAssignVar :: Level -> Level -> State Env Bool
tryAssignVar Level
ai Level
bi = do
  Env
env <- StateT Env Identity Env
forall s (m :: * -> *). MonadState s m => m s
get
  let mb :: Maybe Level
mb = Level -> VarAssign -> Maybe Level
forall a. Level -> IntMap a -> Maybe a
IM.lookup Level
bi (VarAssign -> Maybe Level) -> VarAssign -> Maybe Level
forall a b. (a -> b) -> a -> b
$ Env -> VarAssign
vass Env
env
  if -- already assigned
     | Just Level
x <- Maybe Level
mb
     , Level
x Level -> Level -> Bool
forall a. Eq a => a -> a -> Bool
== Level
ai -> Bool -> State Env Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True

     -- not assigned and levels match
     | Maybe Level
Nothing <- Maybe Level
mb
     , Just Level
lA <- Level -> VarAssign -> Maybe Level
forall a. Level -> IntMap a -> Maybe a
IM.lookup Level
ai (VarAssign -> Maybe Level) -> VarAssign -> Maybe Level
forall a b. (a -> b) -> a -> b
$ Env -> VarAssign
vlA Env
env
     , Just Level
lB <- Level -> VarAssign -> Maybe Level
forall a. Level -> IntMap a -> Maybe a
IM.lookup Level
bi (VarAssign -> Maybe Level) -> VarAssign -> Maybe Level
forall a b. (a -> b) -> a -> b
$ Env -> VarAssign
vlB Env
env
     , Level
lA Level -> Level -> Bool
forall a. Eq a => a -> a -> Bool
== Level
lB
     -> do Env -> StateT Env Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put Env
env { vass :: VarAssign
vass = Level -> Level -> VarAssign -> VarAssign
forall a. Level -> a -> IntMap a -> IntMap a
IM.insert Level
bi Level
ai (VarAssign -> VarAssign) -> VarAssign -> VarAssign
forall a b. (a -> b) -> a -> b
$ Env -> VarAssign
vass Env
env }
           Bool -> State Env Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True

     | Bool
otherwise -> Bool -> State Env Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

-- | Add vars from both sigs to the context, accounting for level
introVars :: [FreeVarIdx]
          -> [FreeVarIdx]
          -> State Env Bool
introVars :: [Level] -> [Level] -> State Env Bool
introVars [] [] = Bool -> State Env Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
introVars [Level]
va [Level]
vb
  | [Level] -> Level
forall (t :: * -> *) a. Foldable t => t a -> Level
length [Level]
va Level -> Level -> Bool
forall a. Eq a => a -> a -> Bool
== [Level] -> Level
forall (t :: * -> *) a. Foldable t => t a -> Level
length [Level]
vb
  = (Bool
True Bool -> StateT Env Identity () -> State Env Bool
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$) (StateT Env Identity () -> State Env Bool)
-> ((Env -> Env) -> StateT Env Identity ())
-> (Env -> Env)
-> State Env Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Env -> Env) -> StateT Env Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' ((Env -> Env) -> State Env Bool) -> (Env -> Env) -> State Env Bool
forall a b. (a -> b) -> a -> b
$ \Env
env ->
      let lvl :: Level
lvl = Env -> Level
level Env
env
       in Env
env { vlA :: VarAssign
vlA = [(Level, Level)] -> VarAssign
forall a. [(Level, a)] -> IntMap a
IM.fromList ([Level] -> [Level] -> [(Level, Level)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Level]
va ([Level] -> [(Level, Level)]) -> [Level] -> [(Level, Level)]
forall a b. (a -> b) -> a -> b
$ Level -> [Level]
forall a. a -> [a]
repeat Level
lvl) VarAssign -> VarAssign -> VarAssign
forall a. Semigroup a => a -> a -> a
<> Env -> VarAssign
vlA Env
env
              , vlB :: VarAssign
vlB = [(Level, Level)] -> VarAssign
forall a. [(Level, a)] -> IntMap a
IM.fromList ([Level] -> [Level] -> [(Level, Level)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Level]
vb ([Level] -> [(Level, Level)]) -> [Level] -> [(Level, Level)]
forall a b. (a -> b) -> a -> b
$ Level -> [Level]
forall a. a -> [a]
repeat Level
lvl) VarAssign -> VarAssign -> VarAssign
forall a. Semigroup a => a -> a -> a
<> Env -> VarAssign
vlB Env
env
              , level :: Level
level = Level
lvl Level -> Level -> Level
forall a. Num a => a -> a -> a
+ Level
1
              }
  | Bool
otherwise = Bool -> State Env Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

-- | Logical conjuction
(/\) :: State env Bool
     -> State env Bool
     -> State env Bool
State env Bool
a /\ :: State env Bool -> State env Bool -> State env Bool
/\ State env Bool
b = do
  Bool
r <- State env Bool
a
  if Bool
r then State env Bool
b else Bool -> State env Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

checkAnd :: [State Env Bool]
         -> State Env Bool
checkAnd :: [State Env Bool] -> State Env Bool
checkAnd = (State Env Bool -> State Env Bool -> State Env Bool)
-> State Env Bool -> [State Env Bool] -> State Env Bool
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' State Env Bool -> State Env Bool -> State Env Bool
forall env. State env Bool -> State env Bool -> State env Bool
(/\) (Bool -> State Env Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True)

-- | Logical disjunction. Discards state if False
(\/) :: State env Bool
     -> State env Bool
     -> State env Bool
State env Bool
a \/ :: State env Bool -> State env Bool -> State env Bool
\/ State env Bool
b = (env -> Identity (Bool, env)) -> State env Bool
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((env -> Identity (Bool, env)) -> State env Bool)
-> (env -> Identity (Bool, env)) -> State env Bool
forall a b. (a -> b) -> a -> b
$ \env
env ->
  let (Bool
ar, env
as) = State env Bool -> env -> (Bool, env)
forall s a. State s a -> s -> (a, s)
runState State env Bool
a env
env
      ~(Bool
br, env
bs) = State env Bool -> env -> (Bool, env)
forall s a. State s a -> s -> (a, s)
runState State env Bool
b env
env
   in if Bool
ar then (Bool, env) -> Identity (Bool, env)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
ar, env
as)
            else if Bool
br then (Bool, env) -> Identity (Bool, env)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
br, env
bs)
                 else (Bool, env) -> Identity (Bool, env)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
False, env
env)

checkOr :: [State env Bool]
        -> State env Bool
checkOr :: [State env Bool] -> State env Bool
checkOr = (State env Bool -> State env Bool -> State env Bool)
-> State env Bool -> [State env Bool] -> State env Bool
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' State env Bool -> State env Bool -> State env Bool
forall env. State env Bool -> State env Bool -> State env Bool
(\/) (Bool -> State env Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False)