{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeOperators #-}
module ASTfunc where

import AST


idxToInt :: Idx env t -> Int
idxToInt Z = 0
idxToInt (S i) = succ (idxToInt i)

data LetBound s env t t' =
    forall env'. LetBound (LHS s t env env') (Tup (V s env') t)

lhsCopy :: Tup s t -> LetBound s env t t
lhsCopy T0 = LetBound (L0 T0) T0
lhsCopy (T1 t) = LetBound (L1 t) (T1 (V t Z))
lhsCopy (T2 t1 t2)
  | LetBound l1 vs1 <- lhsCopy t1
  , LetBound l2 vs2 <- lhsCopy t2
  = LetBound (L2 l1 l2) (T2 (fmapTup (weakenVar (weakenWithLHS l2)) vs1) vs2)

lhstypeOf :: LHS s t env env' -> Tup s t
lhstypeOf (L0 t) = t
lhstypeOf (L1 t) = T1 t
lhstypeOf (L2 l1 l2) = T2 (lhstypeOf l1) (lhstypeOf l2)

rebuildLHS :: LHS s t env env' -> Exists (LHS s t env1)
rebuildLHS (L0 t) = Exists (L0 t)
rebuildLHS (L1 t) = Exists (L1 t)
rebuildLHS (L2 l1 l2)
  | Exists l1' <- rebuildLHS l1
  , Exists l2' <- rebuildLHS l2
  = Exists (L2 l1' l2')

fmapTup :: (forall t'. s t' -> s' t') -> Tup s t -> Tup s' t
fmapTup _ T0 = T0
fmapTup f (T1 x) = T1 (f x)
fmapTup f (T2 t1 t2) = T2 (fmapTup f t1) (fmapTup f t2)

prj :: PVal tenv env -> Idx env t -> Maybe t
prj VZ _ = Nothing
prj (VS _ x) Z = Just x
prj (VS v _) (S i) = prj v i

prjV :: Val env -> Idx env t -> t
prjV v Z = case v of VS _ x -> x
prjV v (S i) = case v of VS v' _ -> prjV v' i

vpush :: LHS s t env env' -> t -> Val env -> Val env'
vpush (L0 _) _ v = v
vpush (L1 _) x v = VS v x
vpush (L2 l1 l2) (x, y) v = vpush l2 y (vpush l1 x v)

reifyNum :: STy t -> NumDict t
reifyNum TInt = NumDict
reifyNum TFloat = NumDict

(>.) :: env2 :> env3 -> env1 :> env2 -> env1 :> env3
Weaken f >. Weaken g = Weaken (f . g)

weakenId :: env :> env
weakenId = Weaken id

weakenSucc :: env :> (t ': env)
weakenSucc = Weaken S

weakenSucc' :: (t ': env) :> env
weakenSucc' = Weaken (\case S i -> i
                            Z -> error "weakenSucc': reference to skipped item")

weakenWithLHS :: LHS s t env env' -> env :> env'
weakenWithLHS (L0 _) = Weaken id
weakenWithLHS (L1 _) = Weaken S
weakenWithLHS (L2 l1 l2) = weakenWithLHS l2 >. weakenWithLHS l1

weakenSkip :: LHS s t env env' -> env :> env2 -> env' :> env2
weakenSkip (L0 _) = id
weakenSkip (L1 _) = (>. weakenSucc')
weakenSkip (L2 l1 l2) = weakenSkip l2 . weakenSkip l1

weakenVar :: env :> env' -> V s env t -> V s env' t
weakenVar w (V t i) = V t (w >:> i)

weakenExpr :: env :> env' -> Expr env t -> Expr env' t
weakenExpr w (Expr e) = Expr (weakenPExpr weakenExpr w e)

weakenPExpr :: (forall env1 env1' t'. env1 :> env1' -> expr env1 t' -> expr env1' t')
            -> env :> env' -> PExpr expr env t -> PExpr expr env' t
weakenPExpr _ _ (Const t x) = Const t x
weakenPExpr f w (Pair e1 e2) = Pair (f w e1) (f w e2)
weakenPExpr _ _ Nil = Nil
weakenPExpr f w (Prim op e) = Prim op (f w e)
weakenPExpr _ w (Var var) = Var (weakenVar w var)
weakenPExpr f w (Let lhs rhs e)
  | Exists lhs' <- rebuildLHS lhs
  = Let lhs' (f w rhs) (f (sinkWithLHS lhs lhs' w) e)

sinkWithLHS :: LHS s t env1 env1' -> LHS s t env2 env2' -> env1 :> env2 -> env1' :> env2'
sinkWithLHS (L0 _) (L0 _) w = w
sinkWithLHS (L1 _) (L1 _) w = Weaken (\case Z -> Z ; S i -> S (w >:> i))
sinkWithLHS (L2 l1 l2) (L2 l1' l2') w = sinkWithLHS l2 l2' (sinkWithLHS l1 l1' w)
sinkWithLHS _ _ _ = error "sinkWithLHS: unequal LHSs"

untupleE :: Tup (Expr env) t -> Expr env t
untupleE = untupleE' Expr

untupleE' :: (forall t'. PExpr expr env t' -> expr env t')
          -> Tup (expr env) t -> expr env t
untupleE' f T0 = f Nil
untupleE' _ (T1 e) = e
untupleE' f (T2 t1 t2) = f (Pair (untupleE' f t1) (untupleE' f t2))

evars :: Tup (V STy env) t -> Expr env t
evars = evars' Expr

evars' :: (forall env' t'. PExpr expr env' t' -> expr env' t')
       -> Tup (V STy env) t -> expr env t
evars' f = untupleE' f . fmapTup (f . Var)

etypeOf :: Expr env t -> Tup STy t
etypeOf (Expr e) = eptypeOf etypeOf e

eptypeOf :: (forall env' t'. expr env' t' -> Tup STy t')
         -> PExpr expr env t -> Tup STy t
eptypeOf _ (Const t _) = T1 t
eptypeOf f (Pair e1 e2) = T2 (f e1) (f e2)
eptypeOf _ Nil = T0
eptypeOf _ (Prim op _) = T1 (opertypeOf op)
eptypeOf _ (Var (V t _)) = T1 t
eptypeOf f (Let _ _ e) = f e

opertypeOf :: Oper (t -> t') -> STy t'
opertypeOf (Add t) = t
opertypeOf (Mul t) = t
opertypeOf Round = TInt

eInto :: (forall env' t'. expr env' t' -> expr' env' t')
      -> PExpr expr env t -> PExpr expr' env t
eInto _ (Const t x) = Const t x
eInto f (Pair e1 e2) = Pair (f e1) (f e2)
eInto _ Nil = Nil
eInto f (Prim op e) = Prim op (f e)
eInto _ (Var v) = Var v
eInto f (Let lhs rhs e) = Let lhs (f rhs) (f e)