Dynamic pattern matching attempt

SimplyDifficult.hs

{-# LANGUAGE TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses, GADTs, FlexibleContexts #-}
module SimplyDifficult where

-- This thing demonstrates dynamic pattern unification,
-- however the examples are too simple to verify that the unifier is
-- working correctly.

-- There were also a thing or two that were way too hard to implement.
-- I've omitted them.

import Data.Set as Set (Set, singleton, unions, union, member, isSubsetOf)
import qualified Data.Set as Set
import Data.Monoid (Any(..), getAny)
import Data.List (unionBy)
import Data.Function (on)
import Data.Foldable (fold)
import qualified Data.Map as Map
import Data.Map (Map)

import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Except
import Control.Monad.Identity
import Control.Monad.Trans.Maybe
import Control.Applicative (Alternative, (<|>), empty, optional, many)
import Debug.Trace

data Trail a = B0 | Trail a :< a
    deriving (Eq,Show)

instance Functor Trail where
    fmap f B0 = B0
    fmap f (xs:<x) = fmap f xs :< f x

instance Foldable Trail where
    foldr f y B0 = y
    foldr f y (xs:<x) = foldr f (f x y) xs

instance Traversable Trail where
    sequenceA B0 = pure B0 
    sequenceA (xs:<x) = pure (:<) <*> sequenceA xs <*> x

(<.>) :: Trail a -> Trail a -> Trail a
a <.> B0 = a
a <.> (xs:<x) = (a <.> xs) :< x

------------------------------------------------------------
-- Terms
------------------------------------------------------------
newtype Bind a = Bind a
    deriving (Eq,Show)

data Term
    = Rigid RName (Trail (Elim Term))
    | Flex Name (Trail (Elim Term))
    | Abs (Bind Term)
    | Pi Term (Bind Term)
    | Sigma Term (Bind Term)
    | Pair Term Term
    | Star
    deriving (Eq,Show)

data Head = Var RName | Meta Name
    deriving (Eq,Show)
data Elim a = App a | Fst | Snd
    deriving (Eq,Show)

data RName
    = Unquoted Int
    | Closed Int
    | Bound Name
    deriving (Eq,Show,Ord)
 
newtype Name = Name Int
    deriving (Eq,Show,Ord)
 
type Type = Term

instance Functor Elim where
    fmap f (App a) = App (f a)
    fmap f Fst = Fst
    fmap f Snd = Snd

instance Foldable Elim where
    foldMap f (App a) = (f a)
    foldMap f Fst = mempty
    foldMap f Snd = mempty
 
------------------------------------------------------------
-- Evaluator
------------------------------------------------------------
data Env = Env {context :: Signature, scope :: [Value]}

data Value
    = VFlex Name (Trail (Elim Value))
    | VRigid RName (Trail (Elim Value))
    | VAbs (Value -> Value)
    | VPi Value (Value -> Value)
    | VSigma Value (Value -> Value)
    | VPair Value Value
    | VStar

extend :: Value -> Env -> Env
extend v env = env { scope = v : scope env }
 
bindEval :: Env -> Bind Term -> (Value -> Value)
bindEval env (Bind term) x = eval (extend x env) term
 
eval :: Env -> Term -> Value
eval env (Rigid h e) = vsp (rigidEval env h) (fmap (fmap (eval env)) e) 
eval env (Flex m e) = vsp (vmeta (context env) m) (fmap (fmap (eval env)) e) 
eval env (Abs e) = VAbs (bindEval env e)
eval env (Pi a b) = VPi (eval env a) (bindEval env b)
eval env (Sigma a b) = VSigma (eval env a) (bindEval env b)
eval env (Pair a b) = VPair (eval env a) (eval env b)
eval env Star = VStar

rigidEval :: Env -> RName -> Value
rigidEval env (Closed i) = scope env !! i
rigidEval env name = vpar (context env) name

vsp :: Value -> Trail (Elim Value) -> Value
vsp head (xs:<App arg) = vsp head xs `vapp` arg
vsp head (xs:<Fst) = vfst (vsp head xs)
vsp head (xs:<Snd) = vsnd (vsp head xs)
vsp head B0 = head

vapp :: Value -> Value -> Value
vapp (VAbs f)       v = f v
vapp (VFlex i e)    v = VFlex i (e:<App v)
vapp (VRigid n e)   v = VRigid n (e:<App v)
vapp _              v = error "runtime type error"
 
vfst :: Value -> Value
vfst (VPair a b)    = a
vfst (VFlex i e)    = VFlex i (e:<Fst)
vfst (VRigid n e)   = VRigid n (e:<Fst)
vfst _              = error "runtime type error"

vsnd :: Value -> Value
vsnd (VPair a b)    = b
vsnd (VFlex i e)    = VFlex i (e:<Snd)
vsnd (VRigid n e)   = VRigid n (e:<Snd)
vsnd _              = error "runtime type error"
 
vmeta :: Signature -> Name -> Value
vmeta ctx i = case Map.lookup i ctx of
    Just (_, DEFN v) -> eval (Env ctx []) v
    Just (_, HOLE) -> VFlex i B0
    Nothing -> VFlex i B0

vpar :: Signature -> RName -> Value
-- vpar ctx (Bound i) = case lookup i ctx of
--     Just (_, DEFN v) -> v
--     Just (_, HOLE) -> VRigid (Bound i) B0
--     Nothing -> VRigid (Bound i) B0
vpar ctx n = VRigid n B0

vunquoted :: Int -> Value
vunquoted i = VRigid (Unquoted i) B0

------------------------------------------------------------
-- Binding
------------------------------------------------------------
class Binding t where
    binding :: Int -> (Int -> RName -> RName) -> t -> t

instance Binding Term where
    binding i f (Rigid h e) = Rigid (f i h) (fmap (fmap (binding i f)) e)
    binding i f (Flex m e) = Flex m (fmap (fmap (binding i f)) e)
    binding i f (Abs e) = Abs (binding i f e)
    binding i f (Pi t t') = Pi (binding i f t) (binding i f t')
    binding i f (Sigma t t') = Sigma (binding i f t) (binding i f t')
    binding i f (Pair a b) = Pair (binding i f a) (binding i f b)
    binding i f Star = Star
 
instance Binding t => Binding (Bind t) where
    binding i f (Bind t) = Bind (binding (i+1) f t)
 
bind :: Binding t => Name -> t -> Bind t
bind k v = Bind (binding 0 help v)
    where help :: Int -> RName -> RName
          help i (Bound k') | (k == k') = (Closed i)
          --help i (Meta k') | (k == k') = Var (Closed i)
          help i head = head

unbind :: Binding t => Name -> Bind t -> t
unbind name (Bind v) = binding 0 help v
    where help :: Int -> RName -> RName
          help i (Closed j) | (i == j) = (Bound name)
          help i head = head

loosen :: Binding t => t -> t
loosen v = binding 0 help v
    where help :: Int -> RName -> RName
          help i (Closed j) | (i <= j) = (Closed (j+1))
          help i head = head

tighten :: Binding t => Int -> t -> t
tighten k v = binding k help v
    where help :: Int -> RName -> RName
          help i (Closed j) | (i < j) = (Closed (j-1))
          help i (Closed j) | (i == j) = error "bug, did you use fv?"
          help i head = head

tighten' :: Binding t => Trail Int -> t -> t
tighten' ts v = binding 0 help v
    where help :: Int -> RName -> RName
          help i (Closed j) | (i <= j) = Closed (j - sum (take' (j-i) ts))
          help i head = head

take' :: Int -> Trail a -> [a]
take' n B0 = []
take' 0 (xs:<x) = []
take' n (xs:<x) = x : take' (n-1) xs

raise v k = binding 0 help v
    where help :: Int -> RName -> RName
          help i (Closed j) | (i <= j) = (Closed (j+k))
          help i head = head

-- This alone should be sufficient for flowing values inside metavars.
remap :: Term -> Trail Term -> Term
remap t xs = abss (binding 0 help t) xs
    where help :: Int -> RName -> RName
          help i (Closed j) | (i <= j) = Closed (locate xs (j-i) 0+i)
          help i head = head

abss :: Term -> Trail Term -> Term
abss term B0 = term
abss term (xs:<_) = abss (Abs (Bind term)) xs

locate :: Trail Term -> Int -> Int -> Int
locate B0 _ _ = error "should have been there"
locate (xs:<Rigid (Closed j) B0) i n | (i == j) = n
locate (xs:<x) i n = locate xs i (n+1)

-- subst :: Binding t => Head -> Bind t -> t
-- subst head (Bind v) = binding 0 help v
--     where help :: Int -> Head -> Head
--           help i (Var (Closed j)) | (i == j) = head
--           help i head = head

-- ------------------------------------------------------------
-- -- Value forcing
-- ------------------------------------------------------------
-- force :: Subs -> Value -> Value
-- force ctx (VFlex m sp) = case lookup ( m) ctx of
--     Just v -> vsp v sp
--     Nothing -> VFlex m sp
-- force ctx v = v
 
------------------------------------------------------------
-- Quotation
------------------------------------------------------------
quote0 :: Value -> Term
quote0 = quote 0

quote :: Int -> Value -> Term
quote i (VFlex j e) = Flex j (fmap (fmap (quote i)) e)
quote i (VRigid h e) = Rigid (varpar i h) (fmap (fmap (quote i)) e)
quote i (VAbs f) = Abs (bindQuote i f)
quote i (VPi a f) = Pi (quote i a) (bindQuote i f)
quote i (VSigma a f) = Sigma (quote i a) (bindQuote i f)
quote i (VPair a b) = Pair (quote i a) (quote i b)
quote i VStar = Star

bindQuote i f = Bind (quote (i+1) (f (vunquoted i)))
 
varpar :: Int -> RName -> RName
varpar i (Unquoted k) = Closed (i - k - 1)
varpar i x            = x

ctxToScope :: Ctx -> [Value]
ctxToScope ctx = snd (help ctx)
    where help :: Ctx -> (Int, [Value])
          help B0 = (0, [])
          help (xs:<x) = let (n, ys) = help xs
                         in (n+1, vunquoted n : ys)

class Nf a where
    nf :: Env -> a -> a

instance Nf Term where
    nf env t = quote (length (scope env)) (eval env t)

instance Nf a => Nf (Bind a) where
    nf env (Bind a) = Bind (nf (extend (vunquoted (length (scope env))) env) a)

instance Nf a => Nf [a] where
    nf env xs = fmap (nf env) xs

-- Some weirdness
-- bonk :: Env
-- bonk = Env Map.empty []
-- 
-- ($$) :: Term -> Term -> Term
-- f $$ a = quote0 $ vapp (eval bonk f) (eval bonk a)
-- 
-- hd :: Term -> Term
-- hd a = quote0 $ vfst (eval bonk a)
-- 
-- tl :: Term -> Term
-- tl a = quote0 $ vsnd (eval bonk a)
 
inst :: Env -> Bind Term -> Term -> Term
inst env a b = quote (length (scope env)) $
                 vapp (eval env (Abs a)) (eval env b)
-- 
-- (%%) :: Term -> Elim Term -> Term
-- t %% e = quote0 $ vsp (eval bonk t) (B0:<(fmap (eval bonk) e))
 
-- Rather than definining functions to determine the free metavariables
-- and variables of terms directly, I use a typeclass to make them
-- available on the whole syntax.
data Flavour = Vars | RigVars | Metas
 
class Occurs t where
    free   :: Flavour -> t -> Set Name

fvs, fvrigs :: Occurs t => t -> Set Name
fvs       = free Vars
fvrigs    = free RigVars
fmvs :: Occurs t => t -> Set Name
fmvs      = free Metas

instance Occurs Term where
    free RigVars  (Rigid (Bound x) e) = singleton x `union` free RigVars e
    free RigVars  (Flex i _) = Set.empty
    free Vars       (Flex _ e)      = free Vars e
    --free RigVars    (Flex _ e)      = free RigVars e
    free Metas      (Flex alpha e)  = singleton alpha `union` free Metas e
    free Vars       (Rigid (Bound x) e) = singleton x `union` free Vars e
    --free RigVars    (Rigid (Bound x) e) = singleton x `union` free RigVars e
    free Metas      (Rigid (Bound _) e) = Set.empty `union` free Metas e
    free l          (Rigid _ e)         = Set.empty `union` free l e
    free l (Abs b) = free l b
    free l (Pi s t) = free l s `union` free l t
    free l (Sigma s t) = free l s `union` free l t
    free l (Pair s t) = free l s `union` free l t
    free l Star = Set.empty
 
instance (Occurs t, Occurs u, Occurs v) => Occurs (t,u,v) where
    free l (a,b,c) = unions [free l a, free l b, free l c]

instance Occurs t => Occurs [t] where
    free l = unions . map (free l)

instance Occurs t => Occurs (Trail t) where
    free l B0 = Set.empty
    free l (xs :< x) = free l xs `union` free l x

instance Occurs t => Occurs (Elim t) where
    free l (App a)       = free l a
    free l Fst           = Set.empty
    free l Snd           = Set.empty

instance Occurs a => Occurs (Bind a) where
    free l (Bind a) = free l a

class Fv a where
    fv :: a -> Set Int
    fvr :: a -> Set Int

instance Fv Term where
    fv (Rigid h e) = fv h `union` fv e
    fv (Flex _ e) = fv e
    fv (Abs b) = fv b
    fv (Pi s t) = fv s `union` fv t
    fv (Sigma s t) = fv s `union` fv t
    fv (Pair s t) = fv s `union` fv t
    fv Star = Set.empty

    fvr (Rigid h e) = fvr h `union` fvr e
    fvr (Flex _ e) = Set.empty
    fvr (Abs b) = fvr b
    fvr (Pi s t) = fvr s `union` fvr t
    fvr (Sigma s t) = fvr s `union` fvr t
    fvr (Pair s t) = fvr s `union` fvr t
    fvr Star = Set.empty

instance Fv RName where
    fv (Closed i) = singleton i
    fv _ = Set.empty

    fvr (Closed i) = singleton i
    fvr _ = Set.empty

instance Fv a => Fv (Trail a) where
    fv B0 = Set.empty
    fv (xs:<x) = fv xs `union` fv x
    fvr B0 = Set.empty
    fvr (xs:<x) = fvr xs `union` fvr x

instance Fv a => Fv (Elim a) where
    fv (App a) = fv a
    fv Fst     = Set.empty
    fv Snd     = Set.empty
    fvr (App a) = fvr a
    fvr Fst     = Set.empty
    fvr Snd     = Set.empty

instance Fv a => Fv (Bind a) where
    fv (Bind a) = (-1) `Set.delete` Set.map (\x -> x-1) (fv a)
    fvr (Bind a) = (-1) `Set.delete` Set.map (\x -> x-1) (fvr a)


data Decl v    = HOLE | DEFN v
    deriving (Eq, Show)
type Entry = (Type, Decl Term)
type Signature = Map Name Entry

-- Metavariables in the signature
--support (xs:<_) = support xs

-- I don't have atom decls, for now, but they would go here.
-- atomDecls :: Signature -> Set Name
-- atomDecls B0 = Set.empty
-- atomDecls (xs:<SA name _) = Set.insert name (atomDecls xs)
-- atomDecls (xs:<_) = atomDecls xs

-- decls :: Signature -> Set Name
-- decls xs = support xs `union` atomDecls xs

metas :: Occurs t => t -> Set Name
metas = fmvs

-- atoms :: Occurs t => t -> Set Name
-- atoms = fvs -- I'm not certain these are atoms.

-- consts :: Occurs t => t -> Set Name
-- consts t = metas t `union` atoms t

type Ctx = Trail Type

data Constraint = C Ctx Ctx Term Term Type Type
    deriving (Eq,Show)

data Solver = Solver {
    signature :: Signature,
    constraints :: [(Status, Constraint)],
    nextVar   :: Int }
    deriving (Eq,Show)

emptySolver :: Int -> Solver
emptySolver = Solver (Map.empty) []

data Status = Active | Blocked
    deriving (Eq,Show)

type Context a
    = ReaderT Ctx (StateT Solver (ExceptT String Identity)) a

fresh :: MonadState Solver m => m Name
fresh = do
    n <- Name <$> gets nextVar
    modify (\ctx -> ctx { nextVar = nextVar ctx + 1 })
    pure n

support :: MonadState Solver m => m (Set Name)
support = help <$> gets signature
    where help :: Signature -> Set Name
          help sig = Set.fromList (Map.keys sig)

runContext :: Ctx -> Solver -> Context a -> Either String (a, Solver)
runContext ctx solver m =
    runIdentity (runExceptT (runStateT (runReaderT m ctx) solver))

index :: Trail a -> Int -> a
index B0 i = error "overshoot"
index (xs:<x) 0 = x
index (xs:<x) n = index xs (n-1)

eq :: (MonadState Solver m, MonadReader Ctx m) => Term -> Type
                                               -> Term -> Type -> m ()
eq s _S t _T = do
    ctx <- ask
    let c = C ctx ctx s t _S _T
    modify (\ctx -> ctx { constraints = (Active, c) : constraints ctx })
    let c' = C ctx ctx _S _T Star Star
    modify (\ctx -> ctx { constraints = (Active, c') : constraints ctx })

lookupVar :: (MonadState Solver m) => Ctx -> RName -> m Type
lookupVar ctx (Closed i) = pure (raise (ctx `index` i) (i+1))

addSignature :: (MonadState Solver m) => Name -> Entry -> m ()
addSignature name se = modify (\sol -> sol { signature = Map.insert name se (signature sol) })

expand :: (MonadReader Ctx m) => Type -> m a -> m a
expand ty m = local (:<ty) m

elaborate :: Term -> Type -> Context Term
elaborate t ty = do
    a0 <- fresh
    ctx <- ask
    addSignature a0 (fromTelescope ctx ty, HOLE)
    let u = Flex a0 (vector 0 ctx)
    case t of
        Rigid name elim -> do
            (f, elim, ty') <- inferRName name >>= inferElims elim
            eq u ty (f elim) ty'
        Flex alpha elim -> do
            (f, elim, ty') <- inferMeta alpha >>= inferElims elim
            eq u ty (f elim) ty'
        Abs (Bind body) -> do
            y1 <- fresh
            y2 <- fresh
            let _Y1 = (Flex y1 (vector 0 ctx))
            let _Y2 = (Flex y2 (vector 0 (ctx:<_Y1)))
            addSignature y1 (fromTelescope ctx Star, HOLE)
            addSignature y2 (fromTelescope (ctx:<_Y1) Star, HOLE)
            u' <- expand _Y1 (elaborate body _Y2)
            eq u ty (Abs (Bind u')) (Pi _Y1 (Bind _Y2))
        Pi a (Bind b) -> do
            s <- elaborate a Star
            t <- expand s (elaborate b Star)
            eq u ty (Pi s (Bind t)) Star
        Sigma a (Bind b) -> do
            s <- elaborate a Star
            t <- expand s (elaborate b Star)
            eq u ty (Sigma s (Bind t)) Star
        Pair a b -> do
            y1 <- fresh
            y2 <- fresh
            let _Y1 = (Flex y1 (vector 0 ctx))
            let _Y2 = (Flex y2 (vector 0 (ctx:<_Y1)))
            addSignature y1 (fromTelescope ctx Star, HOLE)
            addSignature y2 (fromTelescope (ctx:<_Y1) Star, HOLE)
            s <- elaborate a _Y1
            t <- elaborate b (Flex y2 (vector 0 ctx:<App s))
            eq u ty (Pair s t) (Sigma _Y1 (Bind _Y2))
        Star -> eq u Star Star Star
    pure u

type HeadTrail = (Trail (Elim Term) -> Term, Trail (Elim Term), Type)

inferRName :: RName -> Context HeadTrail
inferRName var = do
    ctx <- ask
    ty <- lookupVar ctx var
    pure (Rigid var, B0, ty)

inferMeta :: Name -> Context HeadTrail
inferMeta alpha = do
    isFresh <- (not . member alpha) <$> support
    if isFresh then do
        ctx <- ask
        ty <- fresh
        let _TY = Flex ty (vector 0 ctx)
        addSignature ty (fromTelescope ctx Star, HOLE)
        addSignature alpha (fromTelescope ctx _TY, HOLE)
        pure (Flex alpha, vector 0 ctx, _TY)
    else error "unfresh meta"

inferElims :: Trail (Elim Term) -> HeadTrail -> Context HeadTrail
inferElims B0 ht = pure ht
inferElims (xs:<x) ht = inferElims xs ht >>= inferElim x

inferElim :: Elim Term -> HeadTrail -> Context HeadTrail
inferElim (App x) (f, e, ty) = do
    ctx <- ask
    y1 <- fresh
    y2 <- fresh
    let _Y1 = (Flex y1 (vector 0 ctx))
    let _Y2 = (Flex y2 (vector 0 (ctx:<_Y1)))
    addSignature y1 (fromTelescope ctx Star, HOLE)
    addSignature y2 (fromTelescope (ctx:<_Y1) Star, HOLE)
    u <- elaborate x _Y1
    eq ty Star (Pi _Y1 (Bind _Y2)) Star
    pure (f, e:<App u, Flex y2 (vector 0 ctx:<App u))
inferElim Fst (f, e, ty) = do
    ctx <- ask
    y1 <- fresh
    y2 <- fresh
    let _Y1 = (Flex y1 (vector 0 ctx))
    let _Y2 = (Flex y2 (vector 0 (ctx:<_Y1)))
    addSignature y1 (fromTelescope ctx Star, HOLE)
    addSignature y2 (fromTelescope (ctx:<_Y1) Star, HOLE)
    eq ty Star (Sigma _Y1 (Bind _Y2)) Star
    pure (f, e:<Fst, _Y1)
inferElim Snd (f, e, ty) = do
    ctx <- ask
    y1 <- fresh
    y2 <- fresh
    let _Y1 = (Flex y1 (vector 0 ctx))
    let _Y2 = (Flex y2 (vector 0 (ctx:<_Y1)))
    addSignature y1 (fromTelescope ctx Star, HOLE)
    addSignature y2 (fromTelescope (ctx:<_Y1) Star, HOLE)
    eq ty Star (Sigma _Y1 (Bind _Y2)) Star
    pure (f, e:<Snd, Flex y2 (vector 0 ctx:<App (f (e:<Fst))))

fromTelescope :: Trail Type -> Type -> Type
fromTelescope xs t = foldr (\x y -> Pi x (Bind y)) t xs
--fromTelescope B0 t = t
--fromTelescope (xs:<x) t = fromTelescope xs (Pi x (Bind t))

telescope :: Type -> (Trail Type, Type)
telescope t = telescope' B0 t
    where telescope' :: Trail Type -> Type -> (Trail Type, Type)
          telescope' xs (Pi a (Bind b)) = telescope' (xs :< a) b
          telescope' xs a = (xs, a)

vector :: Int -> Trail a -> Trail (Elim Term)
vector i B0 = B0 
vector i (xs:<x) = vector (i+1) xs :< App (Rigid (Closed  i) B0)

hole :: Ctx -> Type -> Solving Name
hole _Gam _T = do
    alpha <- fresh
    sol <- get
    put sol { signature = Map.insert alpha (fromTelescope _Gam _T, HOLE)
                          (signature sol) }
    pure alpha

define :: Ctx -> Name -> Type -> Term -> Solving ()
define _Gam alpha _S v = do
    sol <- get
    put sol { signature = Map.insert alpha
                          (fromTelescope _Gam _S, DEFN (abss v _Gam))
                          (signature sol) }

normalize :: Signature -> Constraint -> Constraint
normalize sig (C ctx ctx' t t' ty ty') = C
    (nfContext sig ctx)
    (nfContext sig ctx')
    (nf (Env sig (ctxToScope ctx)) t)
    (nf (Env sig (ctxToScope ctx')) t')
    (nf (Env sig (ctxToScope ctx)) ty)
    (nf (Env sig (ctxToScope ctx')) ty')

nfContext :: Signature -> Ctx -> Ctx
nfContext sig B0 = B0
nfContext sig (xs:<x) = let xs' = nfContext sig xs
                        in xs':<nf (Env sig (ctxToScope xs')) x

postpone :: Status -> Constraint -> Solving ()
postpone status constraint = do
  --trace (show (status, constraint)) $
    modify (\sol -> sol { constraints = (status, constraint) : constraints sol })

active :: Constraint -> Solving ()
active = postpone Active

block :: Constraint -> Solving ()
block = postpone Blocked

id_ty :: Term
id_ty = Pi Star (Bind (Pi (Rigid (Closed 0) B0) (Bind (Rigid (Closed 1) B0))))

id_term :: Term
id_term = Abs (Bind (Abs (Bind (Rigid (Closed 0) B0))))

var :: Int -> Term
var i = Rigid (Closed i) B0

ab :: Term -> Term
ab t = Abs (Bind t)

pi' :: Term -> Term -> Term
pi' t u = Pi t (Bind u)

const_ty :: Term
const_ty = Pi Star (Bind (Pi Star (Bind
              (Pi (var 1) (Bind (Pi (var 1) (Bind (var 3))))))))

const_term :: Term
const_term = Abs (Bind (Abs (Bind (Abs (Bind (Abs (Bind (var 1))))))))

num_ty = pi' Star (pi' (pi' (var 0) (var 1)) (pi' (var 1) (var 2)))
zero_term = ab (ab (ab (var 0)))
one_term = ab (ab (ab (Rigid (Closed 1) (B0:<App (var 0)))))
two_term = ab (ab (ab (Rigid (Closed 1) (B0:<App (Rigid (Closed 1) (B0:<App (var 0)))))))

--call_id_term = ab (ab (ab (var 0)))

call_id_ty = (pi' id_ty (pi' Star Star))
call_id_term = (ab (ab (var 0)))
call_id_term2 = (ab (ab (Rigid (Closed 1) (B0:<App Star:<App (var 0)))))

call_id_ty2 = pi' Star (pi' id_ty (pi' (var 1) (var 2)))
call_id_term3 = ab (ab (ab (Rigid (Closed 1) (B0:<App (Flex (Name (-1)) B0):<App (var 0)))))

--call_id_term = ab (ab (ab (Rigid (Closed 1) (B0:<App (Flex (Name (-1)) B0):<App (var 0)))))

runSolver test = do
    (u, solver) <- runContext B0 (emptySolver 0) test
    let solver' = solve solver
    case solver' of
        (Solver sig [] _) -> Right (nf (Env sig []) u, sig)
        (Solver sig cons _) -> Left (show (sig, cons))

runTest1Print :: IO ()
runTest1Print = do
    case runSolver (elaborate const_ty Star) of
        Left error -> putStrLn ("error: " ++ error)
        Right (ty, solver) -> do
            putStrLn ("success: " ++ show ty)
            -- let result = runContext B0 (emptySolver 0) (elaborate id_term ty)
            -- putStrLn (show result)
            case runSolver (elaborate const_term ty) of
                Left error -> putStrLn ("error: " ++ error)
                Right (term, solver) -> do
                    putStrLn ("success: " ++ show term)
    case runSolver (elaborate id_ty Star) of
        Left error -> putStrLn ("error: " ++ error)
        Right (ty, solver) -> do
            putStrLn ("success: " ++ show ty)
            -- let result = runContext B0 (emptySolver 0) (elaborate id_term ty)
            -- putStrLn (show result)
            case runSolver (elaborate id_term ty) of
                Left error -> putStrLn ("error: " ++ error)
                Right (term, solver) -> do
                    putStrLn ("success: " ++ show term)
    case runSolver (elaborate num_ty Star) of
        Left error -> putStrLn ("error: " ++ error)
        Right (ty, solver) -> do
            putStrLn ("success: " ++ show ty)
            case runSolver (elaborate zero_term ty) of
                Left error -> putStrLn ("error: " ++ error)
                Right (term, solver) -> do
                    putStrLn ("success: " ++ show term)
            case runSolver (elaborate one_term ty) of
                Left error -> putStrLn ("error: " ++ error)
                Right (term, solver) -> do
                    putStrLn ("success: " ++ show term)
            case runSolver (elaborate two_term ty) of
                Left error -> putStrLn ("error: " ++ error)
                Right (term, solver) -> do
                    putStrLn ("success: " ++ show term)
    case runSolver (elaborate call_id_ty Star) of
        Left error -> putStrLn ("error: " ++ error)
        Right (ty, solver) -> do
            putStrLn ("success: " ++ show ty)
            case runSolver (elaborate call_id_term ty) of
                Left error -> putStrLn ("error: " ++ error)
                Right (term, solver) -> do
                    putStrLn ("success: " ++ show term)
            case runSolver (elaborate call_id_term2 ty) of
                Left error -> putStrLn ("error: " ++ error)
                Right (term, solver) -> do
                    putStrLn ("success: " ++ show term)
    case runSolver (elaborate call_id_ty2 Star) of
        Left error -> putStrLn ("error: " ++ error)
        Right (ty, solver) -> do
            putStrLn ("success: " ++ show ty)
            case runSolver (elaborate call_id_term3 ty) of
                Left error -> putStrLn ("error: " ++ error)
                Right (term, solver) -> do
                    putStrLn ("success: " ++ show term)

solve :: Solver -> Solver
solve solver = case runIdentity (runStateT (runMaybeT solveSteps) solver) of
               (Nothing, solver') -> solver'
               (Just (), solver') -> solve solver'

type Solving a = MaybeT (StateT Solver Identity) a

solveSteps :: Solving ()
solveSteps = do
    sol <- get
    let cons = constraints sol
    put sol { constraints = [] }
    progress <- solveConstraints cons
    guard progress <|> (gets signature >>= lowering . Map.toList >>= guard)
    
    -- <|> TODO: unblockConstraints by eta-expanding
    -- uninstantiated metavariables

solveConstraints :: [(Status, Constraint)] -> Solving Bool
solveConstraints [] = pure False
solveConstraints ((a,x):xs) = do
    progress <- solveConstraints xs
    sig <- gets signature
    let x' = normalize sig x
    if isActive a || (x' /= x) then do
        optional (refine x')
        pure True
    else do
        block x
        pure progress

isActive :: Status -> Bool
isActive Active = True
isActive Blocked = False

(<||) :: Monad m => m Bool -> m () -> m ()
a <|| b = do x <- a
             unless x b

infixr 5 <||

sym :: Constraint -> Constraint
sym (C ctx ctx' t t' ty ty') = C ctx' ctx t' t ty' ty

refine :: Constraint -> Solving ()
refine (C ctx1 ctx2 t1 t2 ty1 ty2) | (t1 == t2) = pure ()
refine (C ctx1 ctx2 t1 t2 ty1 ty2) = do
        sig <- gets signature
        --trace ("refine " ++ (show (C ctx1 ctx2 t1 t2 ty1 ty2))) $
        --  pure ()
        -- The if case of Definition 2.158 (strongly neutral term) has a
        -- condition on the number of argumetns which may be fullfilled
        -- by n-expanding the terms.
        let t1'' = etaExpandDefHeaded t1
        let t2'' = etaExpandDefHeaded t2
        case (t1'', t2'') of
            (Rigid a e1, Rigid b e2) -> do -- by rule schema 14 (strongly neutral term)
                optional (matchSpine ctx1 a e1 ctx2 b e2)
                block (C ctx1 ctx2 t1'' t2'' ty1 ty2)
            (Flex a e1, Flex b e2) | a == b -> do
                flexFlexSame ctx1 ctx2 a e1 e2 ty1 ty2
            (Flex a e1, Flex b e2) -> do
                let t1''' = etaContractWhnf t1''
                let t2''' = etaContractWhnf t2''
                let q = (C ctx1 ctx2 t1''' t2''' ty1 ty2)
                tryPrune q <|| tryPrune (sym q) <|| tryInvert q <|| tryInvert (sym q) <|| block q
            (Flex a e1, u) -> do
                let t1''' = etaContractWhnf t1''
                let t2''' = etaContractWhnf t2''
                let q = (C ctx1 ctx2 t1''' t2''' ty1 ty2)
                tryPrune q <|| tryInvert q <|| block q
            (u, Flex b e2) -> do
                let t1''' = etaContractWhnf t1''
                let t2''' = etaContractWhnf t2''
                let q = (C ctx2 ctx1 t2''' t1''' ty2 ty1)
                tryPrune q <|| tryInvert q <|| block q
            (_, _) -> do
                (t1''', ty1') <- etaExpand ctx1 t1'' (nf (Env sig (ctxToScope ctx1)) ty1)
                (t2''', ty2') <- etaExpand ctx2 t2'' (nf (Env sig (ctxToScope ctx2)) ty2)
                splitDown ctx1 ctx2 t1''' t2''' ty1' ty2'
                -- by rule schema 11 (λ-abstraction)
                -- by rule schema 12 (pairs)
                -- by rule schema 13 (booleans)
                -- by rule schema 3 (injectivity of ∏)
                -- by rule schema 4 (injectivity of Σ)
                -- by rule schema 5 (bool)
                -- by rule schema 6 set
                --pure [(Never, C ctx1 ctx2 t1''' t2''' ty1' ty2')]

splitDown :: Ctx -> Ctx -> Term -> Term -> Type -> Type -> Solving ()
splitDown ctx ctx' (Pi a (Bind b)) (Pi a' (Bind b')) Star Star = do
    let c1 = C ctx ctx' a a' Star Star
    let c2 = C (ctx:<a) (ctx':<a') b b' Star Star
    active c1
    active c2
splitDown ctx1 ctx2 s t _S _T = block (C ctx1 ctx2 s t _S _T)

tryInvert :: Constraint -> Solving Bool
tryInvert (C ctx1 ctx2 (Flex alpha e1) t2 ty1 ty2) = do
    ty <- getType alpha
    m <- invert alpha ty e1 t2
    case m of
        Nothing ->
            pure False
        Just t2' -> do
            active (C ctx1 ctx2 (Flex alpha e1) t2 ty1 ty2)
            define B0 alpha ty t2'
            pure True

invert :: Name -> Type -> Trail (Elim Term) -> Term -> Solving (Maybe Term)
invert alpha _T e t | alpha `Set.notMember` fmvs t,
                      Just xs <- varList e,
                      fv t `Set.isSubsetOf` fvd xs = do
                          pure (Just (remap t xs))
                    | otherwise = pure Nothing

fvd :: Trail Term -> Set Int
fvd B0 = Set.empty
fvd (xs:<x) = (fvd xs `Set.difference` fv x) `Set.union` (fv x `Set.difference` fvd xs)

getType :: Name -> Solving Type
getType name = do
    sig <- gets signature
    case Map.lookup name sig of
        Just (ty,_) -> pure ty
        Nothing -> error "variable not installed"

tryPrune :: Constraint -> Solving Bool
tryPrune q@(C _ _ (Flex alpha e) t _ _) = do
    u <- pruneTerm (fv e) t
    case u of
        (d:_) ->
            --instantiate d
            --sig <- gets signature
            --let q' = normalize sig q
            --trace "pruned" (active q')
            --pure True
            active q >> instantiate d >> pure True
        [] -> pure False

pruneTerm :: Set Int -> Term -> Solving [Instantiation]
pruneTerm vs (Pi s t) = (++) <$> pruneTerm vs s <*> pruneUnder vs t
pruneTerm vs (Sigma s t) = (++) <$> pruneTerm vs s <*> pruneUnder vs t
pruneTerm vs (Pair s t) = (++) <$> pruneTerm vs s <*> pruneTerm vs t
pruneTerm vs (Abs b) = pruneUnder vs b
pruneTerm vs Star = pure []
pruneTerm vs (Rigid r e) = pruneElims vs e
pruneTerm vs (Flex beta e) = pruneMeta vs beta e

pruneUnder :: Set Int -> Bind Term -> Solving [Instantiation]
pruneUnder vs (Bind term) = pruneTerm (Set.insert 0 (Set.map (+1) vs)) term

pruneElims :: Set Int -> Trail (Elim Term) -> Solving [Instantiation]
pruneElims vs e = fold <$> traverse pruneElim e
    where pruneElim :: Elim Term -> Solving [Instantiation]
          pruneElim (App x) = pruneTerm vs x
          pruneElim Fst = pure []
          pruneElim Snd = pure []

pruneMeta :: Set Int -> Name -> Trail (Elim Term) -> Solving [Instantiation]
pruneMeta vs alpha e = do
    (tel, ty) <- telescope <$> getType alpha
    case pruneVars vs (fv ty) tel e of
        Just dropList | any (/=0) dropList -> do
            let ty' = tighten' dropList ty
            let tel' = tightenCtx dropList tel
            pure [(alpha, fromTelescope tel' ty',
                  \beta -> abss (Flex beta (revec dropList (vector 0 tel))) tel)]
        _ -> pure []

tightenCtx :: Trail Int -> Ctx -> Ctx
tightenCtx _ B0 = B0
tightenCtx (ts:<_) (xs:<x) = tightenCtx ts xs :< tighten' ts x

revec :: Trail Int -> Trail (Elim Term) -> Trail (Elim Term)
revec _ B0 = B0
revec (ts:<0) (xs:<x) = revec ts xs :< x
revec (ts:<_) (xs:<x) = revec ts xs

pruneVars :: Set Int -> Set Int -> Ctx -> Trail (Elim Term) -> Maybe (Trail Int)
pruneVars vs ks B0 B0 = Just B0
pruneVars vs ks (xs :< _S) (e :< App s) =
    if 0 `member` ks
    then do dropList <- pruneVars vs (Set.map (subtract 1) ks `union` fv _S) xs e
            Just (dropList :< 0)
    else case toVar s of
        Just y | y `member` vs -> do
            dropList <- pruneVars vs (Set.map (subtract 1) ks `union` fv _S) xs e
            Just (dropList :< 0)
        _      | not (fvr s `Set.isSubsetOf` vs) -> do
            dropList <- pruneVars vs (Set.map (subtract 1) ks) xs e
            Just (dropList :< 1)
               | otherwise -> Nothing
pruneVars vs ks _ _ = Nothing

-- Similar to prune, so it's down here
flexFlexSame :: Ctx -> Ctx -> Name
             -> Trail (Elim Term) -> Trail (Elim Term)
             -> Type -> Type -> Solving ()
flexFlexSame ctx ctx' alpha e e' ty ty' = do
    (tel, _T) <- telescope <$> getType alpha
    case intersect (fv _T) tel e e' of
        Just dropList | any (/=0) dropList -> let _T' = tighten' dropList _T
                                                  tel' = tightenCtx dropList tel
                                              in instantiate (alpha, fromTelescope tel' _T',
                                                              \beta -> abss (Flex beta (revec dropList (vector 0 tel))) tel)
        _                                  -> block (C ctx ctx' (Flex alpha e) (Flex alpha e') ty ty')

-- Given a telescope and the two evaluation contexts, |intersect|
-- checks the evaluation contexts are lists of variables and produces the
-- telescope on which they agree.
 
intersect :: Set Int -> Trail Type -> Trail (Elim Term) -> Trail (Elim Term) -> Maybe (Trail Int)
intersect ks B0 B0 B0 = pure B0
intersect ks (tel :< _S) (e :< App s) (e' :< App s') = do
    x <- toVar s
    y <- toVar s'
    if x == y then do
        tel <- intersect (Set.map (subtract 1) ks `union` fv _S) tel e e'
        pure (tel :< 0)
    else if 0 `member` ks then empty
         else do tel <- intersect (Set.map (subtract 1) ks) tel e e'
                 pure (tel :< 1)
        
-- For rigid-rigid interaction.
matchSpine :: Ctx -> RName -> Trail (Elim Term)
           -> Ctx -> RName -> Trail (Elim Term)
           -> Solving (Type, Type)
matchSpine ctx x B0 ctx' x' B0
    | x == x' = do
        a <- lookupVar ctx x
        b <- lookupVar ctx' x'
        pure (a, b)
    | otherwise = empty
matchSpine ctx x (e :< App a) ctx' x' (e' :< App a') = do
    (Pi _A _B, Pi _S _T) <- matchSpine ctx x e ctx' x' e'
    active $ C ctx ctx' a a' _A _S
    pure (inst (Env Map.empty (ctxToScope ctx)) _B a,
          inst (Env Map.empty (ctxToScope ctx')) _T a')
matchSpine ctx x (e :< Fst) ctx' x' (e' :< Fst) = do
    (Sigma _A _B, Sigma _S _T) <- matchSpine ctx x e ctx' x' e'
    pure (_A, _S)
matchSpine ctx x (e :< Snd) ctx' x' (e' :< Snd) = do
    (Sigma _A _B, Sigma _S _T) <- matchSpine ctx x e ctx' x' e'
    pure (inst (Env Map.empty (ctxToScope ctx))  _B (Rigid x (e :< Fst)),
          inst (Env Map.empty (ctxToScope ctx')) _T (Rigid x' (e' :< Fst)))
matchSpine _ _ _ _ _ _ = empty

lowering :: [(Name, Entry)] -> Solving Bool
lowering [] = pure False
lowering ((alpha,(ty,HOLE)):xs) = do
    status <- lowering xs
    status' <- lower B0 alpha ty
    pure (status || status')
lowering (_:xs) = lowering xs

-- Given the name and type of a metavariable, |lower| attempts to
-- simplify it by removing $\Sigma$-types, according to the metavariable
-- simplification steps \eqref{eqn:miller:metasimp:sigma} and
-- \eqref{eqn:miller:metasimp:pi-sigma} in
-- Figure~\longref{fig:miller:solve}, as described in
-- Subsection~\longref{subsec:miller:spec:metasimp}.

lower :: Ctx -> Name -> Type -> Solving Bool
lower phi alpha (Sigma _S _T) = do
    beta <- hole phi _S
    let betaTerm = Flex beta (vector 0 phi)
    ceta <- hole phi (inst (Env Map.empty (ctxToScope phi)) _T betaTerm)
    let cetaTerm = Flex ceta (vector 0 phi)
    define phi alpha (Sigma _S _T) (Pair betaTerm cetaTerm)
    pure True
lower phi alpha (Pi _S (Bind _T)) = lower (phi :< _S) alpha _T
lower _ _ _ = pure False

-- > lower :: Telescope -> Nom -> Type -> Contextual Bool
-- > lower _Phi alpha (Sig _S _T) =  hole _Phi _S $ \ s ->
-- >                                 hole _Phi (inst _T s) $ \ t ->
-- >                                 define _Phi alpha (Sig _S _T) (PAIR s t) >>
-- >                                 return True
-- >
-- > lower _Phi alpha (Pi _S _T) = do  x <- fresh (s2n "x")
-- >                                   splitSig B0 x _S >>= maybe
-- >                                       (lower (_Phi :< (x, _S)) alpha (inst _T (var x)))
-- >                                       (\ (y, _A, z, _B, s, (u, v)) ->
-- >                                           hole _Phi (_Pi y _A  (_Pi z _B (inst _T s))) $ \ w ->
-- >                                           define _Phi alpha (Pi _S _T) (lam x (w $$ u $$ v)) >>
-- >                                           return True)      
-- >             
-- > lower _Phi alpha _T = return False
-- 
-- Lowering a metavariable needs to split $\Sigma$-types (possibly
-- underneath a bunch of parameters) into their components.  For example,
-- $[[y : Pi x : X . Sigma z : S . T]]$ splits into
-- $[[y0 : Pi x : X . S]]$ and $[[y1 : Pi x : X . {y0 x/z} T]]$.  Given
-- the name of a variable and its type, |splitSig| attempts to split it,
-- returning fresh variables for the two components of the $\Sigma$-type,
-- an inhabitant of the original type in terms of the new variables and
-- inhabitants of the new types by projecting the original variable.
-- 
-- > splitSig ::  Telescope -> Nom -> Type ->
-- >                  Contextual (Maybe  (Nom, Type, Nom, Type, Tm, (Tm, Tm)))
-- > splitSig _Phi x (Sig _S _T)  = do  y  <- fresh (s2n "y")
-- >                                    z  <- fresh (s2n "z")
-- >                                    return $ Just  (y, _Pis _Phi _S, z, _Pis _Phi (inst _T (var y $*$ _Phi)),
-- >                                                   lams' _Phi (PAIR (var y $*$ _Phi) (var z $*$ _Phi)),
-- >                                                   (lams' _Phi (var x $*$ _Phi %% Hd), 
-- >                                                        lams' _Phi (var x $*$ _Phi %% Tl)))
-- > splitSig _Phi x (Pi _A _B)   = do  a <- fresh (s2n "a")
-- >                                    splitSig (_Phi :< (a, _A)) x (inst _B (var a))
-- > splitSig _ _ _ = return Nothing

type Instantiation = (Name, Type, Name -> Term)

instantiate :: Instantiation -> Solving ()
instantiate d@(alpha, ty, f) = do
    beta <- hole B0 ty
    ty' <- getType alpha
    define B0 alpha ty' (f beta)
    

tZip :: Trail a -> Trail b -> Trail (a, b)
tZip B0 _ = B0
tZip _ B0 = B0
tZip (xs:<x) (ys:<y) = tZip xs ys :< (x, y)

size :: Trail a -> Int
size B0 = 0
size (xs:<_) = size xs + 1

toVar :: Term -> Maybe Int
toVar (Rigid (Closed n) B0) = Just n
toVar _ = Nothing

varList :: Alternative m => Trail (Elim (Term)) -> m (Trail Term)
varList (xs:<App (Flex a B0)) = (:<Flex a B0) <$> varList xs
varList (xs:<App (Rigid n B0)) = (:<Rigid n B0) <$> varList xs
varList B0 = pure B0
varList _ = empty


etaContractWhnf :: Term -> Term
etaContractWhnf t = t
--etaContractWhnf (Abs (Bind (VRigid head (elems:<App (VRigid (Closed 0) B0)))))
--    | not (0 `Set.member` fv (head, elems)) = 

etaExpand :: Ctx -> Term -> Type -> Solving (Term, Type)
etaExpand ctx (Abs x) (Pi a b) = pure (Abs x, Pi a b)
etaExpand ctx (Pair x y) (Sigma a b) = pure (Pair x y, Sigma a b)
etaExpand ctx f (Pi a b) = pure (Abs (Bind expr), Pi a b)
    where expr = quote (length scop) $
                       eval (Env Map.empty scop) (loosen f) `vapp` (vunquoted (length scop))
          scop = (ctxToScope ctx)
etaExpand ctx p (Sigma a b) = pure (Pair e1 e2, Sigma a b)
    where e1 = quote (length scop) $ vfst $ eval (Env Map.empty scop) p
          e2 = quote (length scop) $ vsnd $ eval (Env Map.empty scop) p
          scop = (ctxToScope ctx)
etaExpand ctx p q = pure (p, q)

etaExpandDefHeaded :: Term -> Term
etaExpandDefHeaded t = t