diff options
Diffstat (limited to 'src/Language/AST.hs')
| -rw-r--r-- | src/Language/AST.hs | 21 |
1 files changed, 19 insertions, 2 deletions
diff --git a/src/Language/AST.hs b/src/Language/AST.hs index 3b04bec..f5203e9 100644 --- a/src/Language/AST.hs +++ b/src/Language/AST.hs @@ -30,6 +30,9 @@ data NExpr env t where NEVar :: Lookup name env ~ t => Var name t -> NExpr env t NELet :: Var name a -> NExpr env a -> NExpr ('(name, a) : env) t -> NExpr env t + -- environment management + NEDrop :: SNat i -> NExpr (DropNth i env) t -> NExpr env t + -- base types NEPair :: NExpr env a -> NExpr env b -> NExpr env (TPair a b) NEFst :: NExpr env (TPair a b) -> NExpr env a @@ -41,7 +44,6 @@ data NExpr env t where -- array operations NEConstArr :: Show (ScalRep t) => SNat n -> SScalTy t -> Array n (ScalRep t) -> NExpr env (TArr n (TScal t)) - NEBuild1 :: NExpr env TIx -> Var name TIx -> NExpr ('(name, TIx) : env) t -> NExpr env (TArr (S Z) t) NEBuild :: SNat n -> NExpr env (Tup (Replicate n TIx)) -> Var name (Tup (Replicate n TIx)) -> NExpr ('(name, Tup (Replicate n TIx)) : env) t -> NExpr env (TArr n t) NEFold1Inner :: Var name1 t -> Var name2 t -> NExpr ('(name2, t) : '(name1, t) : env) t -> NExpr env t -> NExpr env (TArr (S n) t) -> NExpr env (TArr n t) NESum1Inner :: ScalIsNumeric t ~ True => NExpr env (TArr (S n) (TScal t)) -> NExpr env (TArr n (TScal t)) @@ -68,6 +70,10 @@ type family Lookup name env where Lookup name ('(name, t) : env) = t Lookup name (_ : env) = Lookup name env +type family DropNth i env where + DropNth Z (_ : env) = env + DropNth (S i) (p : env) = p : DropNth i env + data Var name t = Var (SSymbol name) (STy t) deriving (Show) @@ -135,6 +141,8 @@ fromNamedExpr val = \case \expression to De Bruijn expression" NELet n a b -> ELet ext (go a) (lambda val n b) + NEDrop i e -> weakenExpr (dropNthW i val) (fromNamedExpr (dropNth i val) e) + NEPair a b -> EPair ext (go a) (go b) NEFst e -> EFst ext (go e) NESnd e -> ESnd ext (go e) @@ -144,7 +152,6 @@ fromNamedExpr val = \case NECase e n1 a n2 b -> ECase ext (go e) (lambda val n1 a) (lambda val n2 b) NEConstArr n t x -> EConstArr ext n t x - NEBuild1 a n b -> EBuild1 ext (go a) (lambda val n b) NEBuild k a n b -> EBuild ext k (go a) (lambda val n b) NEFold1Inner n1 n2 a b c -> EFold1Inner ext (lambda2 val n1 n2 a) (go b) (go c) NESum1Inner e -> ESum1Inner ext (go e) @@ -185,3 +192,13 @@ fromNamedExpr val = \case injectWrapLet e (arg `SCons` args) = injectWrapLet (ELet ext (weakenExpr (wSinks args) $ fromNamedExpr val arg) e) args + +dropNth :: SNat i -> NEnv env -> NEnv (DropNth i env) +dropNth SZ (val `NPush` _) = val +dropNth (SS i) (val `NPush` p) = dropNth i val `NPush` p +dropNth _ NTop = error "DropNth: index out of range" + +dropNthW :: SNat i -> NEnv env -> UnName (DropNth i env) :> UnName env +dropNthW SZ (_ `NPush` _) = WSink +dropNthW (SS i) (val `NPush` _) = WCopy (dropNthW i val) +dropNthW _ NTop = error "DropNth: index out of range" |