diff options
Diffstat (limited to 'src/AST.hs')
| -rw-r--r-- | src/AST.hs | 32 |
1 files changed, 29 insertions, 3 deletions
@@ -19,6 +19,7 @@ import Data.Functor.Const import Data.Kind (Type) import Data.Int +import AST.Env import AST.Weaken import Data @@ -90,15 +91,17 @@ data Expr x env t where -- array operations EBuild1 :: x (TArr (S Z) t) -> Expr x env TIx -> Expr x (TIx : env) t -> Expr x env (TArr (S Z) t) - EBuild :: x (TArr n t) -> Vec n (Expr x env TIx) -> Expr x (ConsN n TIx env) t -> Expr x env (TArr n t) + EBuild :: x (TArr n t) -> SNat n -> Expr x env (Tup (Replicate n TIx)) -> Expr x (ConsN n TIx env) t -> Expr x env (TArr n t) EFold1 :: x (TArr n t) -> Expr x (t : t : env) t -> Expr x env (TArr (S n) t) -> Expr x env (TArr n t) EUnit :: x (TArr Z t) -> Expr x env t -> Expr x env (TArr Z t) + EReplicate :: x (TArr (S n) t) -> Expr x env (TArr n t) -> Expr x env (TArr (S n) t) -- TODO: unused -- expression operations EConst :: Show (ScalRep t) => x (TScal t) -> SScalTy t -> ScalRep t -> Expr x env (TScal t) EIdx0 :: x t -> Expr x env (TArr Z t) -> Expr x env t EIdx1 :: x (TArr n t) -> Expr x env (TArr (S n) t) -> Expr x env TIx -> Expr x env (TArr n t) EIdx :: x t -> Expr x env (TArr n t) -> Vec n (Expr x env TIx) -> Expr x env t + EShape :: x (Tup (Replicate n TIx)) -> Expr x env (TArr n t) -> Expr x env (Tup (Replicate n TIx)) EOp :: x t -> SOp a t -> Expr x env a -> Expr x env t -- accumulation effect @@ -114,6 +117,18 @@ type Ex = Expr (Const ()) ext :: Const () a ext = Const () +type family Replicate n x where + Replicate Z x = '[] + Replicate (S n) x = x : Replicate n x + +type family Tup env where + Tup '[] = TNil + Tup (t : ts) = TPair (Tup ts) t + +tTup :: SList STy env -> STy (Tup env) +tTup SNil = STNil +tTup (SCons t ts) = STPair (tTup ts) t + type SOp :: Ty -> Ty -> Type data SOp a t where OAdd :: SScalTy a -> SOp (TPair (TScal a) (TScal a)) (TScal a) @@ -151,14 +166,16 @@ typeOf = \case ECase _ _ a _ -> typeOf a EBuild1 _ _ e -> STArr (SS SZ) (typeOf e) - EBuild _ es e -> STArr (vecLength es) (typeOf e) + EBuild _ n _ e -> STArr n (typeOf e) EFold1 _ _ e | STArr (SS n) t <- typeOf e -> STArr n t EUnit _ e -> STArr SZ (typeOf e) + EReplicate _ e | STArr n t <- typeOf e -> STArr (SS n) t EConst _ t _ -> STScal t EIdx0 _ e | STArr _ t <- typeOf e -> t EIdx1 _ e _ | STArr (SS n) t <- typeOf e -> STArr n t EIdx _ e _ | STArr _ t <- typeOf e -> t + -- EShape _ e | STArr n _ <- typeOf e -> _ EOp _ op _ -> opt2 op EWith e1 e2 -> STPair (typeOf e2) (typeOf e1) @@ -214,9 +231,10 @@ subst' f w = \case EInr x t e -> EInr x t (subst' f w e) ECase x e a b -> ECase x (subst' f w e) (subst' (sinkF f) (WCopy w) a) (subst' (sinkF f) (WCopy w) b) EBuild1 x a b -> EBuild1 x (subst' f w a) (subst' (sinkF f) (WCopy w) b) - EBuild x es e -> EBuild x (fmap (subst' f w) es) (subst' (sinkFN (vecLength es) f) (wcopyN (vecLength es) w) e) + EBuild x n a b -> EBuild x n (subst' f w a) (subst' (sinkFN n f) (wcopyN n w) b) EFold1 x a b -> EFold1 x (subst' (sinkF (sinkF f)) (WCopy (WCopy w)) a) (subst' f w b) EUnit x e -> EUnit x (subst' f w e) + EReplicate x e -> EReplicate x (subst' f w e) EConst x t v -> EConst x t v EIdx0 x e -> EIdx0 x (subst' f w e) EIdx1 x a b -> EIdx1 x (subst' f w a) (subst' f w b) @@ -254,6 +272,11 @@ wpopN :: SNat n -> ConsN n TIx env :> env' -> env :> env' wpopN SZ w = w wpopN (SS n) w = wpopN n (WPop w) +wUndoSubenv :: Subenv env env' -> env' :> env +wUndoSubenv SETop = WId +wUndoSubenv (SEYes sub) = WCopy (wUndoSubenv sub) +wUndoSubenv (SENo sub) = WSink .> wUndoSubenv sub + slistIdx :: SList f list -> Idx list t -> f t slistIdx (SCons x _) IZ = x slistIdx (SCons _ list) (IS i) = slistIdx list i @@ -281,3 +304,6 @@ instance (KnownNat n, KnownTy t) => KnownTy (TAccum n t) where knownTy = STAccum class KnownEnv env where knownEnv :: SList STy env instance KnownEnv '[] where knownEnv = SNil instance (KnownTy t, KnownEnv env) => KnownEnv (t : env) where knownEnv = SCons knownTy knownEnv + +ebuildUp1 :: SNat n -> Ex env (Tup (Replicate n TIx)) -> Ex env TIx -> Ex (TIx : env) (TArr n t) -> Ex env (TArr (S n) t) +ebuildUp1 n sh size f = EBuild ext (SS n) (EPair ext sh size) (error "TODO" f) |