diff options
Diffstat (limited to 'src/AST')
| -rw-r--r-- | src/AST/Accum.hs | 137 | ||||
| -rw-r--r-- | src/AST/Bindings.hs | 75 | ||||
| -rw-r--r-- | src/AST/Count.hs | 170 | ||||
| -rw-r--r-- | src/AST/Env.hs | 85 | ||||
| -rw-r--r-- | src/AST/Pretty.hs | 481 | ||||
| -rw-r--r-- | src/AST/Sparse.hs | 290 | ||||
| -rw-r--r-- | src/AST/Sparse/Types.hs | 107 | ||||
| -rw-r--r-- | src/AST/SplitLets.hs | 154 | ||||
| -rw-r--r-- | src/AST/Types.hs | 205 | ||||
| -rw-r--r-- | src/AST/UnMonoid.hs | 249 | ||||
| -rw-r--r-- | src/AST/Weaken.hs | 132 | ||||
| -rw-r--r-- | src/AST/Weaken/Auto.hs | 176 |
12 files changed, 0 insertions, 2261 deletions
diff --git a/src/AST/Accum.hs b/src/AST/Accum.hs deleted file mode 100644 index 988a450..0000000 --- a/src/AST/Accum.hs +++ /dev/null @@ -1,137 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TypeData #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} -module AST.Accum where - -import AST.Types -import Data - - -data AcPrj - = APHere - | APFst AcPrj - | APSnd AcPrj - | APLeft AcPrj - | APRight AcPrj - | APJust AcPrj - | APArrIdx AcPrj - | APArrSlice Nat - --- | @b@ is a small part of @a@, indicated by the projection @p@. -data SAcPrj (p :: AcPrj) (a :: Ty) (b :: Ty) where - SAPHere :: SAcPrj APHere a a - SAPFst :: SAcPrj p a b -> SAcPrj (APFst p) (TPair a t) b - SAPSnd :: SAcPrj p a b -> SAcPrj (APSnd p) (TPair t a) b - SAPLeft :: SAcPrj p a b -> SAcPrj (APLeft p) (TLEither a t) b - SAPRight :: SAcPrj p a b -> SAcPrj (APRight p) (TLEither t a) b - SAPJust :: SAcPrj p a b -> SAcPrj (APJust p) (TMaybe a) b - SAPArrIdx :: SAcPrj p a b -> SAcPrj (APArrIdx p) (TArr n a) b - -- TODO: - -- SAPArrSlice :: SNat m -> SAcPrj (APArrSlice m) (TArr n t) (TArr (n - m) t) -deriving instance Show (SAcPrj p a b) - -type data AIDense = AID | AIS - -data SAIDense d where - SAID :: SAIDense AID - SAIS :: SAIDense AIS -deriving instance Show (SAIDense d) - -type family AcIdx d p t where - AcIdx d APHere t = TNil - AcIdx AID (APFst p) (TPair a b) = AcIdx AID p a - AcIdx AID (APSnd p) (TPair a b) = AcIdx AID p b - AcIdx AIS (APFst p) (TPair a b) = TPair (AcIdx AIS p a) (ZeroInfo b) - AcIdx AIS (APSnd p) (TPair a b) = TPair (ZeroInfo a) (AcIdx AIS p b) - AcIdx d (APLeft p) (TLEither a b) = AcIdx d p a - AcIdx d (APRight p) (TLEither a b) = AcIdx d p b - AcIdx d (APJust p) (TMaybe a) = AcIdx d p a - AcIdx AID (APArrIdx p) (TArr n a) = - -- (index, recursive info) - TPair (Tup (Replicate n TIx)) (AcIdx AID p a) - AcIdx AIS (APArrIdx p) (TArr n a) = - -- ((index, shape info), recursive info) - TPair (TPair (Tup (Replicate n TIx)) (ZeroInfo (TArr n a))) - (AcIdx AIS p a) - -- AcIdx AID (APArrSlice m) (TArr n a) = - -- -- index - -- Tup (Replicate m TIx) - -- AcIdx AIS (APArrSlice m) (TArr n a) = - -- -- (index, array shape) - -- TPair (Tup (Replicate m TIx)) (Tup (Replicate n TIx)) - -type AcIdxD p t = AcIdx AID p t -type AcIdxS p t = AcIdx AIS p t - -acPrjTy :: SAcPrj p a b -> SMTy a -> SMTy b -acPrjTy SAPHere t = t -acPrjTy (SAPFst prj) (SMTPair t _) = acPrjTy prj t -acPrjTy (SAPSnd prj) (SMTPair _ t) = acPrjTy prj t -acPrjTy (SAPLeft prj) (SMTLEither t _) = acPrjTy prj t -acPrjTy (SAPRight prj) (SMTLEither _ t) = acPrjTy prj t -acPrjTy (SAPJust prj) (SMTMaybe t) = acPrjTy prj t -acPrjTy (SAPArrIdx prj) (SMTArr _ t) = acPrjTy prj t - -type family ZeroInfo t where - ZeroInfo TNil = TNil - ZeroInfo (TPair a b) = TPair (ZeroInfo a) (ZeroInfo b) - ZeroInfo (TLEither a b) = TNil - ZeroInfo (TMaybe a) = TNil - ZeroInfo (TArr n t) = TArr n (ZeroInfo t) - ZeroInfo (TScal t) = TNil - -tZeroInfo :: SMTy t -> STy (ZeroInfo t) -tZeroInfo SMTNil = STNil -tZeroInfo (SMTPair a b) = STPair (tZeroInfo a) (tZeroInfo b) -tZeroInfo (SMTLEither _ _) = STNil -tZeroInfo (SMTMaybe _) = STNil -tZeroInfo (SMTArr n t) = STArr n (tZeroInfo t) -tZeroInfo (SMTScal _) = STNil - --- | Info needed to create a zero-valued deep accumulator for a monoid type. --- Should be constructable from a D1. -type family DeepZeroInfo t where - DeepZeroInfo TNil = TNil - DeepZeroInfo (TPair a b) = TPair (DeepZeroInfo a) (DeepZeroInfo b) - DeepZeroInfo (TLEither a b) = TLEither (DeepZeroInfo a) (DeepZeroInfo b) - DeepZeroInfo (TMaybe a) = TMaybe (DeepZeroInfo a) - DeepZeroInfo (TArr n a) = TArr n (DeepZeroInfo a) - DeepZeroInfo (TScal t) = TNil - -tDeepZeroInfo :: SMTy t -> STy (DeepZeroInfo t) -tDeepZeroInfo SMTNil = STNil -tDeepZeroInfo (SMTPair a b) = STPair (tDeepZeroInfo a) (tDeepZeroInfo b) -tDeepZeroInfo (SMTLEither a b) = STLEither (tDeepZeroInfo a) (tDeepZeroInfo b) -tDeepZeroInfo (SMTMaybe a) = STMaybe (tDeepZeroInfo a) -tDeepZeroInfo (SMTArr n t) = STArr n (tDeepZeroInfo t) -tDeepZeroInfo (SMTScal _) = STNil - --- -- | Additional info needed for accumulation. This is empty unless there is --- -- sparsity in the monoid. --- type family AccumInfo t where --- AccumInfo TNil = TNil --- AccumInfo (TPair a b) = TPair (AccumInfo a) (AccumInfo b) --- AccumInfo (TLEither a b) = TLEither (PrimalInfo a) (PrimalInfo b) --- AccumInfo (TMaybe a) = TMaybe (AccumInfo a) --- AccumInfo (TArr n t) = TArr n (AccumInfo t) --- AccumInfo (TScal t) = TNil - --- type family PrimalInfo t where --- PrimalInfo TNil = TNil --- PrimalInfo (TPair a b) = TPair (PrimalInfo a) (PrimalInfo b) --- PrimalInfo (TLEither a b) = TLEither (PrimalInfo a) (PrimalInfo b) --- PrimalInfo (TMaybe a) = TMaybe (PrimalInfo a) --- PrimalInfo (TArr n t) = TArr n (PrimalInfo t) --- PrimalInfo (TScal t) = TNil - --- tPrimalInfo :: SMTy t -> STy (PrimalInfo t) --- tPrimalInfo SMTNil = STNil --- tPrimalInfo (SMTPair a b) = STPair (tPrimalInfo a) (tPrimalInfo b) --- tPrimalInfo (SMTLEither a b) = STLEither (tPrimalInfo a) (tPrimalInfo b) --- tPrimalInfo (SMTMaybe a) = STMaybe (tPrimalInfo a) --- tPrimalInfo (SMTArr n t) = STArr n (tPrimalInfo t) --- tPrimalInfo (SMTScal _) = STNil diff --git a/src/AST/Bindings.hs b/src/AST/Bindings.hs deleted file mode 100644 index 2310f4b..0000000 --- a/src/AST/Bindings.hs +++ /dev/null @@ -1,75 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE QuantifiedConstraints #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} - --- I want to bring various type variables in scope using type annotations in --- patterns, but I don't want to have to mention all the other type parameters --- of the types in question as well then. Partial type signatures (with '_') are --- useful here. -{-# LANGUAGE PartialTypeSignatures #-} -{-# OPTIONS -Wno-partial-type-signatures #-} -module AST.Bindings where - -import AST -import AST.Env -import Data -import Lemmas - - --- binding lists: a let stack without a body. The stack lives in 'env' and defines 'binds'. -data Bindings f env binds where - BTop :: Bindings f env '[] - BPush :: Bindings f env binds -> (STy t, f (Append binds env) t) -> Bindings f env (t : binds) -deriving instance (forall e t. Show (f e t)) => Show (Bindings f env env') -infixl `BPush` - -mapBindings :: (forall env' t'. f env' t' -> g env' t') - -> Bindings f env binds -> Bindings g env binds -mapBindings _ BTop = BTop -mapBindings f (BPush b (t, e)) = BPush (mapBindings f b) (t, f e) - -weakenBindings :: (forall e1 e2 t. e1 :> e2 -> f e1 t -> f e2 t) - -> env1 :> env2 - -> Bindings f env1 binds - -> (Bindings f env2 binds, Append binds env1 :> Append binds env2) -weakenBindings _ w BTop = (BTop, w) -weakenBindings wf w (BPush b (t, x)) = - let (b', w') = weakenBindings wf w b - in (BPush b' (t, wf w' x), WCopy w') - -weakenOver :: SList STy ts -> env :> env' -> Append ts env :> Append ts env' -weakenOver SNil w = w -weakenOver (SCons _ ts) w = WCopy (weakenOver ts w) - -sinkWithBindings :: forall env' env binds f. Bindings f env binds -> env' :> Append binds env' -sinkWithBindings BTop = WId -sinkWithBindings (BPush b _) = WSink .> sinkWithBindings b - -bconcat :: forall f env binds1 binds2. Bindings f env binds1 -> Bindings f (Append binds1 env) binds2 -> Bindings f env (Append binds2 binds1) -bconcat b1 BTop = b1 -bconcat b1 (BPush (b2 :: Bindings _ (Append binds1 env) binds2C) (t, x)) - | Refl <- lemAppendAssoc @binds2C @binds1 @env - = BPush (bconcat b1 b2) (t, x) - -bindingsBinds :: Bindings f env binds -> SList STy binds -bindingsBinds BTop = SNil -bindingsBinds (BPush binds (t, _)) = SCons t (bindingsBinds binds) - -letBinds :: Bindings Ex env binds -> Ex (Append binds env) t -> Ex env t -letBinds BTop = id -letBinds (BPush b (_, rhs)) = letBinds b . ELet ext rhs - -collectBindings :: SList STy env -> Subenv env env' -> Bindings Ex env env' -collectBindings = \env -> fst . go env WId - where - go :: SList STy env -> env :> env0 -> Subenv env env' -> (Bindings Ex env0 env', env0 :> Append env' env0) - go _ _ SETop = (BTop, WId) - go (ty `SCons` env) w (SEYesR sub) = - let (bs, w') = go env (WPop w) sub - in (BPush bs (ty, EVar ext ty (w' .> w @> IZ)), WSink .> w') - go (_ `SCons` env) w (SENo sub) = go env (WPop w) sub diff --git a/src/AST/Count.hs b/src/AST/Count.hs deleted file mode 100644 index ca4d7ab..0000000 --- a/src/AST/Count.hs +++ /dev/null @@ -1,170 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE DeriveGeneric #-} -{-# LANGUAGE DerivingStrategies #-} -{-# LANGUAGE DerivingVia #-} -{-# LANGUAGE EmptyCase #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE QuantifiedConstraints #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TypeOperators #-} -module AST.Count where - -import Data.Functor.Const -import GHC.Generics (Generic, Generically(..)) - -import AST -import AST.Env -import Data - - -data Count = Zero | One | Many - deriving (Show, Eq, Ord) - -instance Semigroup Count where - Zero <> n = n - n <> Zero = n - _ <> _ = Many -instance Monoid Count where - mempty = Zero - -data Occ = Occ { _occLexical :: Count - , _occRuntime :: Count } - deriving (Eq, Generic) - deriving (Semigroup, Monoid) via Generically Occ - -instance Show Occ where - showsPrec d (Occ l r) = showParen (d > 10) $ - showString "Occ " . showsPrec 11 l . showString " " . showsPrec 11 r - --- | One of the two branches is taken -(<||>) :: Occ -> Occ -> Occ -Occ l1 r1 <||> Occ l2 r2 = Occ (l1 <> l2) (max r1 r2) - --- | This code is executed many times -scaleMany :: Occ -> Occ -scaleMany (Occ l Zero) = Occ l Zero -scaleMany (Occ l _) = Occ l Many - -occCount :: Idx env a -> Expr x env t -> Occ -occCount idx = - getConst . occCountGeneral - (\w i o -> if idx2int i == idx2int (w @> idx) then Const o else mempty) - (\(Const o) -> Const o) - (\(Const o1) (Const o2) -> Const (o1 <||> o2)) - (\(Const o) -> Const (scaleMany o)) - - -data OccEnv env where - OccEnd :: OccEnv env -- not necessarily top! - OccPush :: OccEnv env -> Occ -> OccEnv (t : env) - -instance Semigroup (OccEnv env) where - OccEnd <> e = e - e <> OccEnd = e - OccPush e o <> OccPush e' o' = OccPush (e <> e') (o <> o') - -instance Monoid (OccEnv env) where - mempty = OccEnd - -onehotOccEnv :: Idx env t -> Occ -> OccEnv env -onehotOccEnv IZ v = OccPush OccEnd v -onehotOccEnv (IS i) v = OccPush (onehotOccEnv i v) mempty - -(<||>!) :: OccEnv env -> OccEnv env -> OccEnv env -OccEnd <||>! e = e -e <||>! OccEnd = e -OccPush e o <||>! OccPush e' o' = OccPush (e <||>! e') (o <||> o') - -scaleManyOccEnv :: OccEnv env -> OccEnv env -scaleManyOccEnv OccEnd = OccEnd -scaleManyOccEnv (OccPush e o) = OccPush (scaleManyOccEnv e) (scaleMany o) - -occEnvPop :: OccEnv (t : env) -> OccEnv env -occEnvPop (OccPush o _) = o -occEnvPop OccEnd = OccEnd - -occCountAll :: Expr x env t -> OccEnv env -occCountAll = occCountGeneral (const onehotOccEnv) occEnvPop (<||>!) scaleManyOccEnv - -occCountGeneral :: forall r env t x. - (forall env'. Monoid (r env')) - => (forall env' a. env :> env' -> Idx env' a -> Occ -> r env') -- ^ one-hot - -> (forall env' a. r (a : env') -> r env') -- ^ unpush - -> (forall env'. r env' -> r env' -> r env') -- ^ alternation - -> (forall env'. r env' -> r env') -- ^ scale-many - -> Expr x env t -> r env -occCountGeneral onehot unpush alter many = go WId - where - go :: forall env' t'. Monoid (r env') => env :> env' -> Expr x env' t' -> r env' - go w = \case - EVar _ _ i -> onehot w i (Occ One One) - ELet _ rhs body -> re rhs <> re1 body - EPair _ a b -> re a <> re b - EFst _ e -> re e - ESnd _ e -> re e - ENil _ -> mempty - EInl _ _ e -> re e - EInr _ _ e -> re e - ECase _ e a b -> re e <> (re1 a `alter` re1 b) - ENothing _ _ -> mempty - EJust _ e -> re e - EMaybe _ a b e -> re a <> re1 b <> re e - ELNil _ _ _ -> mempty - ELInl _ _ e -> re e - ELInr _ _ e -> re e - ELCase _ e a b c -> re e <> (re a `alter` re1 b `alter` re1 c) - EConstArr{} -> mempty - EBuild _ _ a b -> re a <> many (re1 b) - EFold1Inner _ _ a b c -> many (unpush (unpush (go (WSink .> WSink .> w) a))) <> re b <> re c - ESum1Inner _ e -> re e - EUnit _ e -> re e - EReplicate1Inner _ a b -> re a <> re b - EMaximum1Inner _ e -> re e - EMinimum1Inner _ e -> re e - EConst{} -> mempty - EIdx0 _ e -> re e - EIdx1 _ a b -> re a <> re b - EIdx _ a b -> re a <> re b - EShape _ e -> re e - EOp _ _ e -> re e - ECustom _ _ _ _ _ _ _ a b -> re a <> re b - ERecompute _ e -> re e - EWith _ _ a b -> re a <> re1 b - EAccum _ _ _ a _ b e -> re a <> re b <> re e - EZero _ _ e -> re e - EDeepZero _ _ e -> re e - EPlus _ _ a b -> re a <> re b - EOneHot _ _ _ a b -> re a <> re b - EError{} -> mempty - where - re :: Monoid (r env') => Expr x env' t'' -> r env' - re = go w - - re1 :: Monoid (r env') => Expr x (a : env') t'' -> r env' - re1 = unpush . go (WSink .> w) - - -deleteUnused :: SList f env -> OccEnv env -> (forall env'. Subenv env env' -> r) -> r -deleteUnused SNil OccEnd k = k SETop -deleteUnused (_ `SCons` env) OccEnd k = - deleteUnused env OccEnd $ \sub -> k (SENo sub) -deleteUnused (_ `SCons` env) (OccPush occenv (Occ _ count)) k = - deleteUnused env occenv $ \sub -> - case count of Zero -> k (SENo sub) - _ -> k (SEYesR sub) - -unsafeWeakenWithSubenv :: Subenv env env' -> Expr x env t -> Expr x env' t -unsafeWeakenWithSubenv = \sub -> - subst (\x t i -> case sinkViaSubenv i sub of - Just i' -> EVar x t i' - Nothing -> error "unsafeWeakenWithSubenv: Index occurred that was subenv'd away") - where - sinkViaSubenv :: Idx env t -> Subenv env env' -> Maybe (Idx env' t) - sinkViaSubenv IZ (SEYesR _) = Just IZ - sinkViaSubenv IZ (SENo _) = Nothing - sinkViaSubenv (IS i) (SEYesR sub) = IS <$> sinkViaSubenv i sub - sinkViaSubenv (IS i) (SENo sub) = sinkViaSubenv i sub diff --git a/src/AST/Env.hs b/src/AST/Env.hs deleted file mode 100644 index 422f0f7..0000000 --- a/src/AST/Env.hs +++ /dev/null @@ -1,85 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE EmptyCase #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PatternSynonyms #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE QuantifiedConstraints #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TypeOperators #-} -module AST.Env where - -import Data.Type.Equality - -import AST.Sparse -import AST.Weaken -import CHAD.Types -import Data - - --- | @env'@ is a subset of @env@: each element of @env@ is either included in --- @env'@ ('SEYes') or not included in @env'@ ('SENo'). -data Subenv' s env env' where - SETop :: Subenv' s '[] '[] - SEYes :: forall t t' env env' s. s t t' -> Subenv' s env env' -> Subenv' s (t : env) (t' : env') - SENo :: forall t env env' s. Subenv' s env env' -> Subenv' s (t : env) env' -deriving instance (forall t t'. Show (s t t')) => Show (Subenv' s env env') - -type Subenv = Subenv' (:~:) -type SubenvS = Subenv' Sparse - -pattern SEYesR :: forall tenv tenv'. () - => forall t env env'. (tenv ~ t : env, tenv' ~ t : env') - => Subenv env env' -> Subenv tenv tenv' -pattern SEYesR s = SEYes Refl s - -{-# COMPLETE SETop, SEYesR, SENo #-} - -subList :: (IsSubType s, IsSubTypeSubject s f) => SList f env -> Subenv' s env env' -> SList f env' -subList SNil SETop = SNil -subList (SCons x xs) (SEYes s sub) = SCons (subtApply s x) (subList xs sub) -subList (SCons _ xs) (SENo sub) = subList xs sub - -subenvAll :: (IsSubType s, IsSubTypeSubject s f) => SList f env -> Subenv' s env env -subenvAll SNil = SETop -subenvAll (SCons t env) = SEYes (subtFull t) (subenvAll env) - -subenvNone :: SList f env -> Subenv' s env '[] -subenvNone SNil = SETop -subenvNone (SCons _ env) = SENo (subenvNone env) - -subenvOnehot :: SList f env -> Idx env t -> s t t' -> Subenv' s env '[t'] -subenvOnehot (SCons _ env) IZ sp = SEYes sp (subenvNone env) -subenvOnehot (SCons _ env) (IS i) sp = SENo (subenvOnehot env i sp) -subenvOnehot SNil i _ = case i of {} - -subenvCompose :: IsSubType s => Subenv' s env1 env2 -> Subenv' s env2 env3 -> Subenv' s env1 env3 -subenvCompose SETop SETop = SETop -subenvCompose (SEYes s1 sub1) (SEYes s2 sub2) = SEYes (subtTrans s1 s2) (subenvCompose sub1 sub2) -subenvCompose (SEYes _ sub1) (SENo sub2) = SENo (subenvCompose sub1 sub2) -subenvCompose (SENo sub1) sub2 = SENo (subenvCompose sub1 sub2) - -subenvConcat :: Subenv' s env1 env1' -> Subenv' s env2 env2' -> Subenv' s (Append env2 env1) (Append env2' env1') -subenvConcat sub1 SETop = sub1 -subenvConcat sub1 (SEYes s sub2) = SEYes s (subenvConcat sub1 sub2) -subenvConcat sub1 (SENo sub2) = SENo (subenvConcat sub1 sub2) - -sinkWithSubenv :: Subenv' s env env' -> env0 :> Append env' env0 -sinkWithSubenv SETop = WId -sinkWithSubenv (SEYes _ sub) = WSink .> sinkWithSubenv sub -sinkWithSubenv (SENo sub) = sinkWithSubenv sub - -wUndoSubenv :: Subenv' (:~:) env env' -> env' :> env -wUndoSubenv SETop = WId -wUndoSubenv (SEYes Refl sub) = WCopy (wUndoSubenv sub) -wUndoSubenv (SENo sub) = WSink .> wUndoSubenv sub - -subenvMap :: (forall a a'. f a -> s a a' -> s' a a') -> SList f env -> Subenv' s env env' -> Subenv' s' env env' -subenvMap _ SNil SETop = SETop -subenvMap f (t `SCons` l) (SEYes s sub) = SEYes (f t s) (subenvMap f l sub) -subenvMap f (_ `SCons` l) (SENo sub) = SENo (subenvMap f l sub) - -subenvD2E :: Subenv env env' -> Subenv (D2E env) (D2E env') -subenvD2E SETop = SETop -subenvD2E (SEYesR sub) = SEYesR (subenvD2E sub) -subenvD2E (SENo sub) = SENo (subenvD2E sub) diff --git a/src/AST/Pretty.hs b/src/AST/Pretty.hs deleted file mode 100644 index fef9686..0000000 --- a/src/AST/Pretty.hs +++ /dev/null @@ -1,481 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE DeriveFunctor #-} -{-# LANGUAGE EmptyCase #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE TypeOperators #-} -module AST.Pretty (pprintExpr, ppExpr, ppSTy, ppSMTy, PrettyX(..)) where - -import Control.Monad (ap) -import Data.List (intersperse, intercalate) -import Data.Functor.Const -import qualified Data.Functor.Product as Product -import Data.String (fromString) -import Prettyprinter -import Prettyprinter.Render.String - -import qualified Data.Text.Lazy as TL -import qualified Prettyprinter.Render.Terminal as PT -import System.Console.ANSI (hSupportsANSI) -import System.IO (stdout) -import System.IO.Unsafe (unsafePerformIO) - -import AST -import AST.Count -import AST.Sparse.Types -import CHAD.Types -import Data - - -class PrettyX x where - prettyX :: x t -> String - - prettyXsuffix :: x t -> String - prettyXsuffix x = "<" ++ prettyX x ++ ">" - -instance PrettyX (Const ()) where - prettyX _ = "" - prettyXsuffix _ = "" - - -type SVal = SList (Const String) - -newtype M a = M { runM :: Int -> (a, Int) } - deriving (Functor) -instance Applicative M where { pure x = M (\i -> (x, i)) ; (<*>) = ap } -instance Monad M where { M f >>= g = M (\i -> let (x, j) = f i in runM (g x) j) } - -genId :: M Int -genId = M (\i -> (i, i + 1)) - -nameBaseForType :: STy t -> String -nameBaseForType STNil = "nil" -nameBaseForType (STPair{}) = "p" -nameBaseForType (STEither{}) = "e" -nameBaseForType (STMaybe{}) = "m" -nameBaseForType (STScal STI32) = "n" -nameBaseForType (STScal STI64) = "n" -nameBaseForType (STArr{}) = "a" -nameBaseForType (STAccum{}) = "ac" -nameBaseForType _ = "x" - -genName' :: String -> M String -genName' prefix = (prefix ++) . show <$> genId - -genNameIfUsedIn' :: String -> STy a -> Idx env a -> Expr x env t -> M String -genNameIfUsedIn' prefix ty idx ex - | occCount idx ex == mempty = case ty of STNil -> return "()" - _ -> return "_" - | otherwise = genName' prefix - -genNameIfUsedIn :: STy a -> Idx env a -> Expr x env t -> M String -genNameIfUsedIn = \t -> genNameIfUsedIn' (nameBaseForType t) t - -pprintExpr :: (KnownEnv env, PrettyX x) => Expr x env t -> IO () -pprintExpr = putStrLn . ppExpr knownEnv - -ppExpr :: PrettyX x => SList STy env -> Expr x env t -> String -ppExpr senv e = render $ fst . flip runM 1 $ do - val <- mkVal senv - e' <- ppExpr' 0 val e - let lam = "λ" ++ intercalate " " (reverse (unSList (\(Product.Pair (Const name) ty) -> "(" ++ name ++ " : " ++ ppSTy 0 ty ++ ")") (slistZip val senv))) ++ "." - return $ group $ flatAlt - (hang 2 $ - ppString lam - <> hardline <> e') - (ppString lam <+> e') - where - mkVal :: SList f env -> M (SVal env) - mkVal SNil = return SNil - mkVal (SCons _ v) = do - val <- mkVal v - name <- genName' "arg" - return (Const name `SCons` val) - -ppExpr' :: PrettyX x => Int -> SVal env -> Expr x env t -> M ADoc -ppExpr' d val expr = case expr of - EVar _ _ i -> return $ ppString (getConst (slistIdx val i)) <> ppX expr - - e@ELet{} -> ppExprLet d val e - - EPair _ a b -> do - a' <- ppExpr' 0 val a - b' <- ppExpr' 0 val b - return $ group $ flatAlt (align $ ppString "(" <> a' <> hardline <> ppString "," <> b' <> ppString ")" <> ppX expr) - (ppString "(" <> a' <> ppString "," <+> b' <> ppString ")" <> ppX expr) - - EFst _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppString "fst" <> ppX expr <+> e' - - ESnd _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppString "snd" <> ppX expr <+> e' - - ENil _ -> return $ ppString "()" - - EInl _ _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppString "Inl" <> ppX expr <+> e' - - EInr _ _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppString "Inr" <> ppX expr <+> e' - - ECase _ e a b -> do - e' <- ppExpr' 0 val e - let STEither t1 t2 = typeOf e - name1 <- genNameIfUsedIn t1 IZ a - a' <- ppExpr' 0 (Const name1 `SCons` val) a - name2 <- genNameIfUsedIn t2 IZ b - b' <- ppExpr' 0 (Const name2 `SCons` val) b - return $ ppParen (d > 0) $ - hang 2 $ - annotate AKey (ppString "case") <> ppX expr <+> e' <+> annotate AKey (ppString "of") - <> hardline <> ppString "Inl" <+> ppString name1 <+> ppString "->" <+> a' - <> hardline <> ppString "Inr" <+> ppString name2 <+> ppString "->" <+> b' - - ENothing _ _ -> return $ ppString "Nothing" - - EJust _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppString "Just" <> ppX expr <+> e' - - EMaybe _ a b e -> do - let STMaybe t = typeOf e - e' <- ppExpr' 0 val e - a' <- ppExpr' 0 val a - name <- genNameIfUsedIn t IZ b - b' <- ppExpr' 0 (Const name `SCons` val) b - return $ ppParen (d > 0) $ - align $ - group (flatAlt - (annotate AKey (ppString "case") <> ppX expr <+> e' - <> hardline <> annotate AKey (ppString "of")) - (annotate AKey (ppString "case") <> ppX expr <+> e' <+> annotate AKey (ppString "of"))) - <> hardline - <> indent 2 - (ppString "Nothing" <+> ppString "->" <+> a' - <> hardline <> ppString "Just" <+> ppString name <+> ppString "->" <+> b') - - ELNil _ _ _ -> return (ppString "LNil") - - ELInl _ _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppString "LInl" <> ppX expr <+> e' - - ELInr _ _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppString "LInr" <> ppX expr <+> e' - - ELCase _ e a b c -> do - e' <- ppExpr' 0 val e - let STLEither t1 t2 = typeOf e - a' <- ppExpr' 11 val a - name1 <- genNameIfUsedIn t1 IZ b - b' <- ppExpr' 0 (Const name1 `SCons` val) b - name2 <- genNameIfUsedIn t2 IZ c - c' <- ppExpr' 0 (Const name2 `SCons` val) c - return $ ppParen (d > 0) $ - hang 2 $ - annotate AKey (ppString "lcase") <> ppX expr <+> e' <+> annotate AKey (ppString "of") - <> hardline <> ppString "LNil" <+> ppString "->" <+> a' - <> hardline <> ppString "LInl" <+> ppString name1 <+> ppString "->" <+> b' - <> hardline <> ppString "LInr" <+> ppString name2 <+> ppString "->" <+> c' - - EConstArr _ _ ty v - | Dict <- scalRepIsShow ty -> return $ ppString (showsPrec d v "") <> ppX expr - - EBuild _ n a b -> do - a' <- ppExpr' 11 val a - name <- genNameIfUsedIn' "i" (tTup (sreplicate n tIx)) IZ b - e' <- ppExpr' 0 (Const name `SCons` val) b - let primName = ppString ("build" ++ intSubscript (fromSNat n)) - return $ ppParen (d > 0) $ - group $ flatAlt - (hang 2 $ - annotate AHighlight primName <> ppX expr <+> a' - <+> ppString "$" <+> ppString "\\" <> ppString name <+> ppString "->" - <> hardline <> e') - (ppApp (annotate AHighlight primName <> ppX expr) [a', ppLam [ppString name] e']) - - EFold1Inner _ cm a b c -> do - name1 <- genNameIfUsedIn (typeOf a) (IS IZ) a - name2 <- genNameIfUsedIn (typeOf a) IZ a - a' <- ppExpr' 0 (Const name2 `SCons` Const name1 `SCons` val) a - b' <- ppExpr' 11 val b - c' <- ppExpr' 11 val c - let opname = case cm of Commut -> "fold1i(C)" - Noncommut -> "fold1i" - return $ ppParen (d > 10) $ - ppApp (annotate AHighlight (ppString opname) <> ppX expr) [ppLam [ppString name1, ppString name2] a', b', c'] - - ESum1Inner _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppString "sum1i" <> ppX expr <+> e' - - EUnit _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppString "unit" <> ppX expr <+> e' - - EReplicate1Inner _ a b -> do - a' <- ppExpr' 11 val a - b' <- ppExpr' 11 val b - return $ ppParen (d > 10) $ ppApp (ppString "replicate1i" <> ppX expr) [a', b'] - - EMaximum1Inner _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppString "maximum1i" <> ppX expr <+> e' - - EMinimum1Inner _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppString "minimum1i" <> ppX expr <+> e' - - EConst _ ty v - | Dict <- scalRepIsShow ty -> return $ ppString (showsPrec d v "") <> ppX expr - - EIdx0 _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppString "idx0" <> ppX expr <+> e' - - EIdx1 _ a b -> do - a' <- ppExpr' 9 val a - b' <- ppExpr' 9 val b - return $ ppParen (d > 8) $ a' <+> ppString ".!" <> ppX expr <+> b' - - EIdx _ a b -> do - a' <- ppExpr' 9 val a - b' <- ppExpr' 10 val b - return $ ppParen (d > 8) $ - a' <+> ppString "!" <> ppX expr <+> b' - - EShape _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppString "shape" <> ppX expr <+> e' - - EOp _ op (EPair _ a b) - | (Infix, ops) <- operator op -> do - a' <- ppExpr' 9 val a - b' <- ppExpr' 9 val b - return $ ppParen (d > 8) $ a' <+> ppString ops <> ppX expr <+> b' - - EOp _ op e -> do - e' <- ppExpr' 11 val e - let ops = case operator op of - (Infix, s) -> "(" ++ s ++ ")" - (Prefix, s) -> s - return $ ppParen (d > 10) $ ppString ops <> ppX expr <+> e' - - ECustom _ t1 t2 t3 a b c e1 e2 -> do - en1 <- genNameIfUsedIn t1 (IS IZ) a - en2 <- genNameIfUsedIn t2 IZ a - pn1 <- genNameIfUsedIn (d1 t1) (IS IZ) b - pn2 <- genNameIfUsedIn (d1 t2) IZ b - dn1 <- genNameIfUsedIn' "tape" t3 (IS IZ) c - dn2 <- genNameIfUsedIn' "d" (d2 (typeOf a)) IZ c - a' <- ppExpr' 11 (Const pn2 `SCons` Const pn1 `SCons` SNil) a - b' <- ppExpr' 11 (Const pn2 `SCons` Const pn1 `SCons` SNil) b - c' <- ppExpr' 11 (Const dn2 `SCons` Const dn1 `SCons` SNil) c - e1' <- ppExpr' 11 val e1 - e2' <- ppExpr' 11 val e2 - return $ ppParen (d > 10) $ - ppApp (ppString "custom" <> ppX expr) - [ppLam [ppString en1, ppString en2] a' - ,ppLam [ppString pn1, ppString pn2] b' - ,ppLam [ppString dn1, ppString dn2] c' - ,e1' - ,e2'] - - ERecompute _ e -> do - e' <- ppExpr' 11 val e - return $ ppParen (d > 10) $ ppApp (ppString "recompute" <> ppX expr) [e'] - - EWith _ t e1 e2 -> do - e1' <- ppExpr' 11 val e1 - name <- genNameIfUsedIn' "ac" (STAccum t) IZ e2 - e2' <- ppExpr' 0 (Const name `SCons` val) e2 - return $ ppParen (d > 0) $ - group $ flatAlt - (hang 2 $ - annotate AWith (ppString "with") <> ppX expr <+> e1' - <+> ppString "$" <+> ppString "\\" <> ppString name <+> ppString "->" - <> hardline <> e2') - (ppApp (annotate AWith (ppString "with") <> ppX expr) [e1', ppLam [ppString name] e2']) - - EAccum _ t prj e1 sp e2 e3 -> do - e1' <- ppExpr' 11 val e1 - e2' <- ppExpr' 11 val e2 - e3' <- ppExpr' 11 val e3 - return $ ppParen (d > 10) $ - ppApp (annotate AMonoid (ppString "accum") <> ppX expr <+> ppString "@" <> ppSMTy' 11 (applySparse sp (acPrjTy prj t))) - [ppString (ppAcPrj t prj), ppString (ppSparse (acPrjTy prj t) sp), e1', e2', e3'] - - EZero _ t e1 -> do - e1' <- ppExpr' 11 val e1 - return $ ppParen (d > 0) $ - annotate AMonoid (ppString "zero") <> ppX expr <+> ppString "@" <> ppSMTy' 11 t <+> e1' - - EDeepZero _ t e1 -> do - e1' <- ppExpr' 11 val e1 - return $ ppParen (d > 0) $ - annotate AMonoid (ppString "deepzero") <> ppX expr <+> ppString "@" <> ppSMTy' 11 t <+> e1' - - EPlus _ t a b -> do - a' <- ppExpr' 11 val a - b' <- ppExpr' 11 val b - return $ ppParen (d > 10) $ - ppApp (annotate AMonoid (ppString "plus") <> ppX expr <+> ppString "@" <> ppSMTy' 11 t) [a', b'] - - EOneHot _ t prj a b -> do - a' <- ppExpr' 11 val a - b' <- ppExpr' 11 val b - return $ ppParen (d > 10) $ - ppApp (annotate AMonoid (ppString "onehot") <> ppX expr <+> ppString "@" <> ppSMTy' 11 (acPrjTy prj t)) [ppString (ppAcPrj t prj), a', b'] - - EError _ _ s -> return $ ppParen (d > 10) $ ppString "error" <> ppX expr <+> ppString (show s) - -ppExprLet :: PrettyX x => Int -> SVal env -> Expr x env t -> M ADoc -ppExprLet d val etop = do - let collect :: PrettyX x => SVal env -> Expr x env t -> M ([(String, Occ, ADoc)], ADoc) - collect val' (ELet _ rhs body) = do - let occ = occCount IZ body - name <- genNameIfUsedIn (typeOf rhs) IZ body - rhs' <- ppExpr' 0 val' rhs - (binds, core) <- collect (Const name `SCons` val') body - return ((name, occ, rhs') : binds, core) - collect val' e = ([],) <$> ppExpr' 0 val' e - - (binds, core) <- collect val etop - - return $ ppParen (d > 0) $ - align $ - annotate AKey (ppString "let") - <+> align (mconcat $ intersperse hardline $ - map (\(name, _occ, rhs) -> - ppString (name ++ {- " (" ++ show _occ ++ ")" ++ -} " = ") <> rhs) - binds) - <> hardline <> annotate AKey (ppString "in") <+> core - -ppApp :: ADoc -> [ADoc] -> ADoc -ppApp fun args = group $ fun <+> align (sep args) - -ppLam :: [ADoc] -> ADoc -> ADoc -ppLam args body = ppString "(" <> hang 2 (ppString "\\" <> sep (args ++ [ppString "->"]) - <> softline <> body <> ppString ")") - -ppAcPrj :: SMTy a -> SAcPrj p a b -> String -ppAcPrj _ SAPHere = "." -ppAcPrj (SMTPair t _) (SAPFst prj) = "(" ++ ppAcPrj t prj ++ ",)" -ppAcPrj (SMTPair _ t) (SAPSnd prj) = "(," ++ ppAcPrj t prj ++ ")" -ppAcPrj (SMTLEither t _) (SAPLeft prj) = "(" ++ ppAcPrj t prj ++ "|)" -ppAcPrj (SMTLEither _ t) (SAPRight prj) = "(|" ++ ppAcPrj t prj ++ ")" -ppAcPrj (SMTMaybe t) (SAPJust prj) = "J" ++ ppAcPrj t prj -ppAcPrj (SMTArr n t) (SAPArrIdx prj) = "[" ++ ppAcPrj t prj ++ "]" ++ intSubscript (fromSNat n) - -ppSparse :: SMTy a -> Sparse a b -> String -ppSparse t sp | Just Refl <- isDense t sp = "D" -ppSparse _ SpAbsent = "A" -ppSparse t (SpSparse s) = "S" ++ ppSparse t s -ppSparse (SMTPair t1 t2) (SpPair s1 s2) = "(" ++ ppSparse t1 s1 ++ "," ++ ppSparse t2 s2 ++ ")" -ppSparse (SMTLEither t1 t2) (SpLEither s1 s2) = "(" ++ ppSparse t1 s1 ++ "|" ++ ppSparse t2 s2 ++ ")" -ppSparse (SMTMaybe t) (SpMaybe s) = "M" ++ ppSparse t s -ppSparse (SMTArr _ t) (SpArr s) = "A" ++ ppSparse t s -ppSparse (SMTScal _) SpScal = "." - -ppX :: PrettyX x => Expr x env t -> ADoc -ppX expr = annotate AExt $ ppString $ prettyXsuffix (extOf expr) - -data Fixity = Prefix | Infix - deriving (Show) - -operator :: SOp a t -> (Fixity, String) -operator OAdd{} = (Infix, "+") -operator OMul{} = (Infix, "*") -operator ONeg{} = (Prefix, "negate") -operator OLt{} = (Infix, "<") -operator OLe{} = (Infix, "<=") -operator OEq{} = (Infix, "==") -operator ONot = (Prefix, "not") -operator OAnd = (Infix, "&&") -operator OOr = (Infix, "||") -operator OIf = (Prefix, "ifB") -operator ORound64 = (Prefix, "round") -operator OToFl64 = (Prefix, "toFl64") -operator ORecip{} = (Prefix, "recip") -operator OExp{} = (Prefix, "exp") -operator OLog{} = (Prefix, "log") -operator OIDiv{} = (Infix, "`div`") -operator OMod{} = (Infix, "`mod`") - -ppSTy :: Int -> STy t -> String -ppSTy d ty = render $ ppSTy' d ty - -ppSTy' :: Int -> STy t -> Doc q -ppSTy' _ STNil = ppString "1" -ppSTy' d (STPair a b) = ppParen (d > 7) $ ppSTy' 8 a <> ppString " * " <> ppSTy' 8 b -ppSTy' d (STEither a b) = ppParen (d > 6) $ ppSTy' 7 a <> ppString " + " <> ppSTy' 7 b -ppSTy' d (STLEither a b) = ppParen (d > 6) $ ppSTy' 7 a <> ppString " ⊕ " <> ppSTy' 7 b -ppSTy' d (STMaybe t) = ppParen (d > 10) $ ppString "Maybe " <> ppSTy' 11 t -ppSTy' d (STArr n t) = ppParen (d > 10) $ - ppString "Arr " <> ppString (show (fromSNat n)) <> ppString " " <> ppSTy' 11 t -ppSTy' _ (STScal sty) = ppString $ case sty of - STI32 -> "i32" - STI64 -> "i64" - STF32 -> "f32" - STF64 -> "f64" - STBool -> "bool" -ppSTy' d (STAccum t) = ppParen (d > 10) $ ppString "Accum " <> ppSMTy' 11 t - -ppSMTy :: Int -> SMTy t -> String -ppSMTy d ty = render $ ppSMTy' d ty - -ppSMTy' :: Int -> SMTy t -> Doc q -ppSMTy' _ SMTNil = ppString "1" -ppSMTy' d (SMTPair a b) = ppParen (d > 7) $ ppSMTy' 8 a <> ppString " * " <> ppSMTy' 8 b -ppSMTy' d (SMTLEither a b) = ppParen (d > 6) $ ppSMTy' 7 a <> ppString " ⊕ " <> ppSMTy' 7 b -ppSMTy' d (SMTMaybe t) = ppParen (d > 10) $ ppString "Maybe " <> ppSMTy' 11 t -ppSMTy' d (SMTArr n t) = ppParen (d > 10) $ - ppString "Arr " <> ppString (show (fromSNat n)) <> ppString " " <> ppSMTy' 11 t -ppSMTy' _ (SMTScal sty) = ppString $ case sty of - STI32 -> "i32" - STI64 -> "i64" - STF32 -> "f32" - STF64 -> "f64" - -ppString :: String -> Doc x -ppString = fromString - -ppParen :: Bool -> Doc x -> Doc x -ppParen True = parens -ppParen False = id - -intSubscript :: Int -> String -intSubscript = \case 0 -> "₀" - n | n < 0 -> '₋' : go (-n) "" - | otherwise -> go n "" - where go 0 suff = suff - go n suff = let (q, r) = n `quotRem` 10 - in go q ("₀₁₂₃₄₅₆₇₈₉" !! r : suff) - -data Annot = AKey | AWith | AHighlight | AMonoid | AExt - deriving (Show) - -annotToANSI :: Annot -> PT.AnsiStyle -annotToANSI AKey = PT.bold -annotToANSI AWith = PT.color PT.Red <> PT.underlined -annotToANSI AHighlight = PT.color PT.Blue -annotToANSI AMonoid = PT.color PT.Green -annotToANSI AExt = PT.colorDull PT.White - -type ADoc = Doc Annot - -render :: Doc Annot -> String -render = - (if stdoutTTY then TL.unpack . PT.renderLazy . reAnnotateS annotToANSI - else renderString) - . layoutPretty LayoutOptions { layoutPageWidth = AvailablePerLine 120 1.0 } - where - stdoutTTY = unsafePerformIO $ hSupportsANSI stdout diff --git a/src/AST/Sparse.hs b/src/AST/Sparse.hs deleted file mode 100644 index 93258b7..0000000 --- a/src/AST/Sparse.hs +++ /dev/null @@ -1,290 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE ImpredicativeTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE RankNTypes #-} - -{-# OPTIONS_GHC -fmax-pmcheck-models=80 #-} -module AST.Sparse (module AST.Sparse, module AST.Sparse.Types) where - -import Data.Type.Equality - -import AST -import AST.Sparse.Types -import Data (SBool(..)) - - -sparsePlus :: SMTy t -> Sparse t t' -> Ex env t' -> Ex env t' -> Ex env t' -sparsePlus _ SpAbsent e1 e2 = use e1 $ use e2 $ ENil ext -sparsePlus t sp e1 e2 | Just Refl <- isDense t sp = EPlus ext t e1 e2 -sparsePlus t (SpSparse sp) e1 e2 = sparsePlus (SMTMaybe t) (SpMaybe sp) e1 e2 -- heh -sparsePlus (SMTPair t1 t2) (SpPair sp1 sp2) e1 e2 = - eunPair e1 $ \w1 e1a e1b -> - eunPair (weakenExpr w1 e2) $ \w2 e2a e2b -> - EPair ext (sparsePlus t1 sp1 (weakenExpr w2 e1a) e2a) - (sparsePlus t2 sp2 (weakenExpr w2 e1b) e2b) -sparsePlus (SMTLEither t1 t2) (SpLEither sp1 sp2) e1 e2 = - elet e2 $ - elcase (weakenExpr WSink e1) - (evar IZ) - (elcase (evar (IS IZ)) - (ELInl ext (applySparse sp2 (fromSMTy t2)) (evar IZ)) - (ELInl ext (applySparse sp2 (fromSMTy t2)) (sparsePlus t1 sp1 (evar (IS IZ)) (evar IZ))) - (EError ext (fromSMTy (applySparse (SpLEither sp1 sp2) (SMTLEither t1 t2))) "splus ll+lr")) - (elcase (evar (IS IZ)) - (ELInr ext (applySparse sp1 (fromSMTy t1)) (evar IZ)) - (EError ext (fromSMTy (applySparse (SpLEither sp1 sp2) (SMTLEither t1 t2))) "splus lr+ll") - (ELInr ext (applySparse sp1 (fromSMTy t1)) (sparsePlus t2 sp2 (evar (IS IZ)) (evar IZ)))) -sparsePlus (SMTMaybe t) (SpMaybe sp) e1 e2 = - elet e2 $ - emaybe (weakenExpr WSink e1) - (evar IZ) - (emaybe (evar (IS IZ)) - (EJust ext (evar IZ)) - (EJust ext (sparsePlus t sp (evar (IS IZ)) (evar IZ)))) -sparsePlus (SMTArr _ t) (SpArr sp) e1 e2 = ezipWith (sparsePlus t sp (evar (IS IZ)) (evar IZ)) e1 e2 -sparsePlus t@SMTScal{} SpScal e1 e2 = EPlus ext t e1 e2 - - -cheapZero :: SMTy t -> Maybe (forall env. Ex env t) -cheapZero SMTNil = Just (ENil ext) -cheapZero (SMTPair t1 t2) - | Just e1 <- cheapZero t1 - , Just e2 <- cheapZero t2 - = Just (EPair ext e1 e2) - | otherwise - = Nothing -cheapZero (SMTLEither t1 t2) = Just (ELNil ext (fromSMTy t1) (fromSMTy t2)) -cheapZero (SMTMaybe t) = Just (ENothing ext (fromSMTy t)) -cheapZero SMTArr{} = Nothing -cheapZero (SMTScal t) = case t of - STI32 -> Just (EConst ext t 0) - STI64 -> Just (EConst ext t 0) - STF32 -> Just (EConst ext t 0.0) - STF64 -> Just (EConst ext t 0.0) - - -data Injection sp a b where - -- | 'Inj' is purposefully also allowed when @sp@ is @False@ so that - -- 'sparsePlusS' can provide injections even if the caller doesn't require - -- them. This simplifies the sparsePlusS code. - Inj :: (forall e. Ex e a -> Ex e b) -> Injection sp a b - Noinj :: Injection False a b - -withInj :: Injection sp a b -> ((forall e. Ex e a -> Ex e b) -> (forall e'. Ex e' a' -> Ex e' b')) -> Injection sp a' b' -withInj (Inj f) k = Inj (k f) -withInj Noinj _ = Noinj - -withInj2 :: Injection sp a1 b1 -> Injection sp a2 b2 - -> ((forall e. Ex e a1 -> Ex e b1) - -> (forall e. Ex e a2 -> Ex e b2) - -> (forall e'. Ex e' a' -> Ex e' b')) - -> Injection sp a' b' -withInj2 (Inj f) (Inj g) k = Inj (k f g) -withInj2 Noinj _ _ = Noinj -withInj2 _ Noinj _ = Noinj - -use :: Ex env a -> Ex env b -> Ex env b -use a b = elet a $ weakenExpr WSink b - --- | This function produces quadratically-sized code in the presence of nested --- dynamic sparsity. TODO can this be improved? -sparsePlusS - :: SBool inj1 -> SBool inj2 - -> SMTy t -> Sparse t t1 -> Sparse t t2 - -> (forall t3. Sparse t t3 - -> Injection inj1 t1 t3 -- only available if first injection is requested (second argument may be absent) - -> Injection inj2 t2 t3 -- only available if second injection is requested (first argument may be absent) - -> (forall e. Ex e t1 -> Ex e t2 -> Ex e t3) - -> r) - -> r --- nil override (but don't destroy effects!) -sparsePlusS _ _ SMTNil _ _ k = - k SpAbsent (Inj $ \a -> use a $ ENil ext) (Inj $ \b -> use b $ ENil ext) (\a b -> use a $ use b $ ENil ext) - --- simplifications -sparsePlusS req1 req2 t (SpSparse SpAbsent) sp2 k = - sparsePlusS req1 req2 t SpAbsent sp2 $ \sp3 minj1 minj2 plus -> - k sp3 (withInj minj1 $ \inj1 -> \a -> use a $ inj1 (ENil ext)) minj2 (\a b -> use a $ plus (ENil ext) b) -sparsePlusS req1 req2 t sp1 (SpSparse SpAbsent) k = - sparsePlusS req1 req2 t sp1 SpAbsent $ \sp3 minj1 minj2 plus -> - k sp3 minj1 (withInj minj2 $ \inj2 -> \b -> use b $ inj2 (ENil ext)) (\a b -> use b $ plus a (ENil ext)) - -sparsePlusS req1 req2 t (SpSparse (SpSparse sp1)) sp2 k = - let ta = applySparse sp1 (fromSMTy t) in - sparsePlusS req1 req2 t (SpSparse sp1) sp2 $ \sp3 minj1 minj2 plus -> - k sp3 - (withInj minj1 $ \inj1 -> \a -> inj1 (emaybe a (ENothing ext ta) (EVar ext (STMaybe ta) IZ))) - minj2 - (\a b -> plus (emaybe a (ENothing ext ta) (EVar ext (STMaybe ta) IZ)) b) -sparsePlusS req1 req2 t sp1 (SpSparse (SpSparse sp2)) k = - let tb = applySparse sp2 (fromSMTy t) in - sparsePlusS req1 req2 t sp1 (SpSparse sp2) $ \sp3 minj1 minj2 plus -> - k sp3 - minj1 - (withInj minj2 $ \inj2 -> \b -> inj2 (emaybe b (ENothing ext tb) (EVar ext (STMaybe tb) IZ))) - (\a b -> plus a (emaybe b (ENothing ext tb) (EVar ext (STMaybe tb) IZ))) - -sparsePlusS req1 req2 t (SpSparse (SpLEither sp1a sp1b)) sp2 k = - let STLEither ta tb = applySparse (SpLEither sp1a sp1b) (fromSMTy t) in - sparsePlusS req1 req2 t (SpLEither sp1a sp1b) sp2 $ \sp3 minj1 minj2 plus -> - k sp3 - (withInj minj1 $ \inj1 -> \a -> inj1 (emaybe a (ELNil ext ta tb) (EVar ext (STLEither ta tb) IZ))) - minj2 - (\a b -> plus (emaybe a (ELNil ext ta tb) (EVar ext (STLEither ta tb) IZ)) b) -sparsePlusS req1 req2 t sp1 (SpSparse (SpLEither sp2a sp2b)) k = - let STLEither ta tb = applySparse (SpLEither sp2a sp2b) (fromSMTy t) in - sparsePlusS req1 req2 t sp1 (SpLEither sp2a sp2b) $ \sp3 minj1 minj2 plus -> - k sp3 - minj1 - (withInj minj2 $ \inj2 -> \b -> inj2 (emaybe b (ELNil ext ta tb) (EVar ext (STLEither ta tb) IZ))) - (\a b -> plus a (emaybe b (ELNil ext ta tb) (EVar ext (STLEither ta tb) IZ))) - -sparsePlusS req1 req2 t (SpSparse (SpMaybe sp1)) sp2 k = - let STMaybe ta = applySparse (SpMaybe sp1) (fromSMTy t) in - sparsePlusS req1 req2 t (SpMaybe sp1) sp2 $ \sp3 minj1 minj2 plus -> - k sp3 - (withInj minj1 $ \inj1 -> \a -> inj1 (emaybe a (ENothing ext ta) (evar IZ))) - minj2 - (\a b -> plus (emaybe a (ENothing ext ta) (EVar ext (STMaybe ta) IZ)) b) -sparsePlusS req1 req2 t sp1 (SpSparse (SpMaybe sp2)) k = - let STMaybe tb = applySparse (SpMaybe sp2) (fromSMTy t) in - sparsePlusS req1 req2 t sp1 (SpMaybe sp2) $ \sp3 minj1 minj2 plus -> - k sp3 - minj1 - (withInj minj2 $ \inj2 -> \b -> inj2 (emaybe b (ENothing ext tb) (evar IZ))) - (\a b -> plus a (emaybe b (ENothing ext tb) (EVar ext (STMaybe tb) IZ))) -sparsePlusS req1 req2 t (SpMaybe (SpSparse sp1)) sp2 k = sparsePlusS req1 req2 t (SpSparse (SpMaybe sp1)) sp2 k -sparsePlusS req1 req2 t sp1 (SpMaybe (SpSparse sp2)) k = sparsePlusS req1 req2 t sp1 (SpSparse (SpMaybe sp2)) k - --- TODO: sparse of Just is just Maybe - --- dense plus -sparsePlusS _ _ t sp1 sp2 k - | Just Refl <- isDense t sp1 - , Just Refl <- isDense t sp2 - = k (spDense t) (Inj id) (Inj id) (\a b -> EPlus ext t a b) - --- handle absents -sparsePlusS SF _ _ SpAbsent sp2 k = k sp2 Noinj (Inj id) (\a b -> use a $ b) -sparsePlusS ST _ t SpAbsent sp2 k - | Just zero2 <- cheapZero (applySparse sp2 t) = - k sp2 (Inj $ \a -> use a $ zero2) (Inj id) (\a b -> use a $ b) - | otherwise = - k (SpSparse sp2) (Inj $ \a -> use a $ ENothing ext (applySparse sp2 (fromSMTy t))) (Inj $ EJust ext) (\a b -> use a $ EJust ext b) - -sparsePlusS _ SF _ sp1 SpAbsent k = k sp1 (Inj id) Noinj (\a b -> use b $ a) -sparsePlusS _ ST t sp1 SpAbsent k - | Just zero1 <- cheapZero (applySparse sp1 t) = - k sp1 (Inj id) (Inj $ \b -> use b $ zero1) (\a b -> use b $ a) - | otherwise = - k (SpSparse sp1) (Inj $ EJust ext) (Inj $ \b -> use b $ ENothing ext (applySparse sp1 (fromSMTy t))) (\a b -> use b $ EJust ext a) - --- double sparse yields sparse -sparsePlusS _ _ t (SpSparse sp1) (SpSparse sp2) k = - sparsePlusS ST ST t sp1 sp2 $ \sp3 (Inj inj1) (Inj inj2) plus -> - k (SpSparse sp3) - (Inj $ \a -> emaybe a (ENothing ext (applySparse sp3 (fromSMTy t))) (EJust ext (inj1 (evar IZ)))) - (Inj $ \b -> emaybe b (ENothing ext (applySparse sp3 (fromSMTy t))) (EJust ext (inj2 (evar IZ)))) - (\a b -> - elet b $ - emaybe (weakenExpr WSink a) - (emaybe (evar IZ) - (ENothing ext (applySparse sp3 (fromSMTy t))) - (EJust ext (inj2 (evar IZ)))) - (emaybe (evar (IS IZ)) - (EJust ext (inj1 (evar IZ))) - (EJust ext (plus (evar (IS IZ)) (evar IZ))))) - --- single sparse can yield non-sparse if the other argument is always present -sparsePlusS SF _ t (SpSparse sp1) sp2 k = - sparsePlusS SF ST t sp1 sp2 $ \sp3 _ (Inj inj2) plus -> - k sp3 Noinj (Inj inj2) - (\a b -> - elet b $ - emaybe (weakenExpr WSink a) - (inj2 (evar IZ)) - (plus (evar IZ) (evar (IS IZ)))) -sparsePlusS ST _ t (SpSparse sp1) sp2 k = - sparsePlusS ST ST t sp1 sp2 $ \sp3 (Inj inj1) (Inj inj2) plus -> - k (SpSparse sp3) - (Inj $ \a -> emaybe a (ENothing ext (applySparse sp3 (fromSMTy t))) (EJust ext (inj1 (evar IZ)))) - (Inj $ \b -> EJust ext (inj2 b)) - (\a b -> - elet b $ - emaybe (weakenExpr WSink a) - (EJust ext (inj2 (evar IZ))) - (EJust ext (plus (evar IZ) (evar (IS IZ))))) -sparsePlusS req1 req2 t sp1 (SpSparse sp2) k = - sparsePlusS req2 req1 t (SpSparse sp2) sp1 $ \sp3 inj1 inj2 plus -> - k sp3 inj2 inj1 (flip plus) - --- products -sparsePlusS req1 req2 (SMTPair ta tb) (SpPair sp1a sp1b) (SpPair sp2a sp2b) k = - sparsePlusS req1 req2 ta sp1a sp2a $ \sp3a minj13a minj23a plusa -> - sparsePlusS req1 req2 tb sp1b sp2b $ \sp3b minj13b minj23b plusb -> - k (SpPair sp3a sp3b) - (withInj2 minj13a minj13b $ \inj13a inj13b -> - \x1 -> eunPair x1 $ \_ x1a x1b -> EPair ext (inj13a x1a) (inj13b x1b)) - (withInj2 minj23a minj23b $ \inj23a inj23b -> - \x2 -> eunPair x2 $ \_ x2a x2b -> EPair ext (inj23a x2a) (inj23b x2b)) - (\x1 x2 -> - eunPair x1 $ \w1 x1a x1b -> - eunPair (weakenExpr w1 x2) $ \w2 x2a x2b -> - EPair ext (plusa (weakenExpr w2 x1a) x2a) (plusb (weakenExpr w2 x1b) x2b)) - --- coproducts -sparsePlusS _ _ (SMTLEither ta tb) (SpLEither sp1a sp1b) (SpLEither sp2a sp2b) k = - sparsePlusS ST ST ta sp1a sp2a $ \(sp3a :: Sparse _t3 t3a) (Inj inj13a) (Inj inj23a) plusa -> - sparsePlusS ST ST tb sp1b sp2b $ \(sp3b :: Sparse _t3' t3b) (Inj inj13b) (Inj inj23b) plusb -> - let nil :: Ex e (TLEither t3a t3b) ; nil = ELNil ext (applySparse sp3a (fromSMTy ta)) (applySparse sp3b (fromSMTy tb)) - inl :: Ex e t3a -> Ex e (TLEither t3a t3b) ; inl = ELInl ext (applySparse sp3b (fromSMTy tb)) - inr :: Ex e t3b -> Ex e (TLEither t3a t3b) ; inr = ELInr ext (applySparse sp3a (fromSMTy ta)) - in - k (SpLEither sp3a sp3b) - (Inj $ \x1 -> elcase x1 nil (inl (inj13a (evar IZ))) (inr (inj13b (evar IZ)))) - (Inj $ \x2 -> elcase x2 nil (inl (inj23a (evar IZ))) (inr (inj23b (evar IZ)))) - (\x1 x2 -> - elet x2 $ - elcase (weakenExpr WSink x1) - (elcase (evar IZ) - nil - (inl (inj23a (evar IZ))) - (inr (inj23b (evar IZ)))) - (elcase (evar (IS IZ)) - (inl (inj13a (evar IZ))) - (inl (plusa (evar (IS IZ)) (evar IZ))) - (EError ext (applySparse (SpLEither sp3a sp3b) (fromSMTy (SMTLEither ta tb))) "plusS ll+lr")) - (elcase (evar (IS IZ)) - (inr (inj13b (evar IZ))) - (EError ext (applySparse (SpLEither sp3a sp3b) (fromSMTy (SMTLEither ta tb))) "plusS lr+ll") - (inr (plusb (evar (IS IZ)) (evar IZ))))) - --- maybe -sparsePlusS _ _ (SMTMaybe t) (SpMaybe sp1) (SpMaybe sp2) k = - sparsePlusS ST ST t sp1 sp2 $ \sp3 (Inj inj1) (Inj inj2) plus -> - k (SpMaybe sp3) - (Inj $ \a -> emaybe a (ENothing ext (applySparse sp3 (fromSMTy t))) (EJust ext (inj1 (evar IZ)))) - (Inj $ \b -> emaybe b (ENothing ext (applySparse sp3 (fromSMTy t))) (EJust ext (inj2 (evar IZ)))) - (\a b -> - elet b $ - emaybe (weakenExpr WSink a) - (emaybe (evar IZ) - (ENothing ext (applySparse sp3 (fromSMTy t))) - (EJust ext (inj2 (evar IZ)))) - (emaybe (evar (IS IZ)) - (EJust ext (inj1 (evar IZ))) - (EJust ext (plus (evar (IS IZ)) (evar IZ))))) - --- dense array cotangents simply recurse -sparsePlusS req1 req2 (SMTArr _ t) (SpArr sp1) (SpArr sp2) k = - sparsePlusS req1 req2 t sp1 sp2 $ \sp3 minj1 minj2 plus -> - k (SpArr sp3) - (withInj minj1 $ \inj1 -> emap (inj1 (EVar ext (applySparse sp1 (fromSMTy t)) IZ))) - (withInj minj2 $ \inj2 -> emap (inj2 (EVar ext (applySparse sp2 (fromSMTy t)) IZ))) - (ezipWith (plus (EVar ext (applySparse sp1 (fromSMTy t)) (IS IZ)) - (EVar ext (applySparse sp2 (fromSMTy t)) IZ))) - --- scalars -sparsePlusS _ _ (SMTScal t) SpScal SpScal k = k SpScal (Inj id) (Inj id) (EPlus ext (SMTScal t)) diff --git a/src/AST/Sparse/Types.hs b/src/AST/Sparse/Types.hs deleted file mode 100644 index 10cac4e..0000000 --- a/src/AST/Sparse/Types.hs +++ /dev/null @@ -1,107 +0,0 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -module AST.Sparse.Types where - -import AST.Types - -import Data.Kind (Type, Constraint) -import Data.Type.Equality - - -data Sparse t t' where - SpSparse :: Sparse t t' -> Sparse t (TMaybe t') - SpAbsent :: Sparse t TNil - - SpPair :: Sparse a a' -> Sparse b b' -> Sparse (TPair a b) (TPair a' b') - SpLEither :: Sparse a a' -> Sparse b b' -> Sparse (TLEither a b) (TLEither a' b') - SpMaybe :: Sparse t t' -> Sparse (TMaybe t) (TMaybe t') - SpArr :: Sparse t t' -> Sparse (TArr n t) (TArr n t') - SpScal :: Sparse (TScal t) (TScal t) -deriving instance Show (Sparse t t') - -class ApplySparse f where - applySparse :: Sparse t t' -> f t -> f t' - -instance ApplySparse STy where - applySparse (SpSparse s) t = STMaybe (applySparse s t) - applySparse SpAbsent _ = STNil - applySparse (SpPair s1 s2) (STPair t1 t2) = STPair (applySparse s1 t1) (applySparse s2 t2) - applySparse (SpLEither s1 s2) (STLEither t1 t2) = STLEither (applySparse s1 t1) (applySparse s2 t2) - applySparse (SpMaybe s) (STMaybe t) = STMaybe (applySparse s t) - applySparse (SpArr s) (STArr n t) = STArr n (applySparse s t) - applySparse SpScal t = t - -instance ApplySparse SMTy where - applySparse (SpSparse s) t = SMTMaybe (applySparse s t) - applySparse SpAbsent _ = SMTNil - applySparse (SpPair s1 s2) (SMTPair t1 t2) = SMTPair (applySparse s1 t1) (applySparse s2 t2) - applySparse (SpLEither s1 s2) (SMTLEither t1 t2) = SMTLEither (applySparse s1 t1) (applySparse s2 t2) - applySparse (SpMaybe s) (SMTMaybe t) = SMTMaybe (applySparse s t) - applySparse (SpArr s) (SMTArr n t) = SMTArr n (applySparse s t) - applySparse SpScal t = t - - -class IsSubType s where - type IsSubTypeSubject (s :: k -> k -> Type) (f :: k -> Type) :: Constraint - subtApply :: IsSubTypeSubject s f => s t t' -> f t -> f t' - subtTrans :: s a b -> s b c -> s a c - subtFull :: IsSubTypeSubject s f => f t -> s t t - -instance IsSubType (:~:) where - type IsSubTypeSubject (:~:) f = () - subtApply = gcastWith - subtTrans = trans - subtFull _ = Refl - -instance IsSubType Sparse where - type IsSubTypeSubject Sparse f = f ~ SMTy - subtApply = applySparse - - subtTrans s1 (SpSparse s2) = SpSparse (subtTrans s1 s2) - subtTrans _ SpAbsent = SpAbsent - subtTrans (SpPair s1a s1b) (SpPair s2a s2b) = SpPair (subtTrans s1a s2a) (subtTrans s1b s2b) - subtTrans (SpLEither s1a s1b) (SpLEither s2a s2b) = SpLEither (subtTrans s1a s2a) (subtTrans s1b s2b) - subtTrans (SpSparse s1) (SpMaybe s2) = SpSparse (subtTrans s1 s2) - subtTrans (SpMaybe s1) (SpMaybe s2) = SpMaybe (subtTrans s1 s2) - subtTrans (SpArr s1) (SpArr s2) = SpArr (subtTrans s1 s2) - subtTrans SpScal SpScal = SpScal - - subtFull = spDense - -spDense :: SMTy t -> Sparse t t -spDense SMTNil = SpAbsent -spDense (SMTPair t1 t2) = SpPair (spDense t1) (spDense t2) -spDense (SMTLEither t1 t2) = SpLEither (spDense t1) (spDense t2) -spDense (SMTMaybe t) = SpMaybe (spDense t) -spDense (SMTArr _ t) = SpArr (spDense t) -spDense (SMTScal _) = SpScal - -isDense :: SMTy t -> Sparse t t' -> Maybe (t :~: t') -isDense SMTNil SpAbsent = Just Refl -isDense _ SpSparse{} = Nothing -isDense _ SpAbsent = Nothing -isDense (SMTPair t1 t2) (SpPair s1 s2) - | Just Refl <- isDense t1 s1, Just Refl <- isDense t2 s2 = Just Refl - | otherwise = Nothing -isDense (SMTLEither t1 t2) (SpLEither s1 s2) - | Just Refl <- isDense t1 s1, Just Refl <- isDense t2 s2 = Just Refl - | otherwise = Nothing -isDense (SMTMaybe t) (SpMaybe s) - | Just Refl <- isDense t s = Just Refl - | otherwise = Nothing -isDense (SMTArr _ t) (SpArr s) - | Just Refl <- isDense t s = Just Refl - | otherwise = Nothing -isDense (SMTScal _) SpScal = Just Refl - -isAbsent :: Sparse t t' -> Bool -isAbsent (SpSparse s) = isAbsent s -isAbsent SpAbsent = True -isAbsent (SpPair s1 s2) = isAbsent s1 && isAbsent s2 -isAbsent (SpLEither s1 s2) = isAbsent s1 && isAbsent s2 -isAbsent (SpMaybe s) = isAbsent s -isAbsent (SpArr s) = isAbsent s -isAbsent SpScal = False diff --git a/src/AST/SplitLets.hs b/src/AST/SplitLets.hs deleted file mode 100644 index dcaf82f..0000000 --- a/src/AST/SplitLets.hs +++ /dev/null @@ -1,154 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UndecidableInstances #-} -module AST.SplitLets (splitLets) where - -import Data.Type.Equality - -import AST -import AST.Bindings -import Lemmas - - -splitLets :: Ex env t -> Ex env t -splitLets = splitLets' (\t i w -> EVar ext t (w @> i)) - -splitLets' :: (forall a env2. STy a -> Idx env a -> env' :> env2 -> Ex env2 a) -> Ex env t -> Ex env' t -splitLets' = \sub -> \case - EVar _ t i -> sub t i WId - ELet _ (rhs :: Ex env t1) body -> ELet ext (splitLets' sub rhs) (split1 sub (typeOf rhs) body) - ECase x e a b -> - let STEither t1 t2 = typeOf e - in ECase x (splitLets' sub e) (split1 sub t1 a) (split1 sub t2 b) - EMaybe x a b e -> - let STMaybe t1 = typeOf e - in EMaybe x (splitLets' sub a) (split1 sub t1 b) (splitLets' sub e) - ELCase x e a b c -> - let STLEither t1 t2 = typeOf e - in ELCase x (splitLets' sub e) (splitLets' sub a) (split1 sub t1 b) (split1 sub t2 c) - EFold1Inner x cm a b c -> - let STArr _ t1 = typeOf c - in EFold1Inner x cm (split2 sub t1 t1 a) (splitLets' sub b) (splitLets' sub c) - - EPair x a b -> EPair x (splitLets' sub a) (splitLets' sub b) - EFst x e -> EFst x (splitLets' sub e) - ESnd x e -> ESnd x (splitLets' sub e) - ENil x -> ENil x - EInl x t e -> EInl x t (splitLets' sub e) - EInr x t e -> EInr x t (splitLets' sub e) - ENothing x t -> ENothing x t - EJust x e -> EJust x (splitLets' sub e) - ELNil x t1 t2 -> ELNil x t1 t2 - ELInl x t e -> ELInl x t (splitLets' sub e) - ELInr x t e -> ELInr x t (splitLets' sub e) - EConstArr x n t a -> EConstArr x n t a - EBuild x n a b -> EBuild x n (splitLets' sub a) (splitLets' (sinkF sub) b) - ESum1Inner x e -> ESum1Inner x (splitLets' sub e) - EUnit x e -> EUnit x (splitLets' sub e) - EReplicate1Inner x a b -> EReplicate1Inner x (splitLets' sub a) (splitLets' sub b) - EMaximum1Inner x e -> EMaximum1Inner x (splitLets' sub e) - EMinimum1Inner x e -> EMinimum1Inner x (splitLets' sub e) - EConst x t v -> EConst x t v - EIdx0 x e -> EIdx0 x (splitLets' sub e) - EIdx1 x a b -> EIdx1 x (splitLets' sub a) (splitLets' sub b) - EIdx x e es -> EIdx x (splitLets' sub e) (splitLets' sub es) - EShape x e -> EShape x (splitLets' sub e) - EOp x op e -> EOp x op (splitLets' sub e) - ECustom x s t p a b c e1 e2 -> ECustom x s t p a b c (splitLets' sub e1) (splitLets' sub e2) - ERecompute x e -> ERecompute x (splitLets' sub e) - EWith x t e1 e2 -> EWith x t (splitLets' sub e1) (splitLets' (sinkF sub) e2) - EAccum x t p e1 sp e2 e3 -> EAccum x t p (splitLets' sub e1) sp (splitLets' sub e2) (splitLets' sub e3) - EZero x t ezi -> EZero x t (splitLets' sub ezi) - EDeepZero x t ezi -> EDeepZero x t (splitLets' sub ezi) - EPlus x t a b -> EPlus x t (splitLets' sub a) (splitLets' sub b) - EOneHot x t p a b -> EOneHot x t p (splitLets' sub a) (splitLets' sub b) - EError x t s -> EError x t s - where - sinkF :: (forall a env2. STy a -> Idx env a -> env' :> env2 -> Ex env2 a) - -> STy t -> Idx (b : env) t -> (b : env') :> env3 -> Ex env3 t - sinkF _ t IZ w = EVar ext t (w @> IZ) - sinkF f t (IS i) w = f t i (w .> WSink) - - split1 :: (forall a env2. STy a -> Idx env a -> env' :> env2 -> Ex env2 a) - -> STy bind -> Ex (bind : env) t -> Ex (bind : env') t - split1 sub (tbind :: STy bind) body = - let (ptrs, bs) = split tbind - in letBinds bs $ - splitLets' (\cases _ IZ w -> subPointers ptrs w - t (IS i) w -> sub t i (WPop @bind (wPops (bindingsBinds bs) w))) - body - - split2 :: forall bind1 bind2 env' env t. - (forall a env2. STy a -> Idx env a -> env' :> env2 -> Ex env2 a) - -> STy bind1 -> STy bind2 -> Ex (bind2 : bind1 : env) t -> Ex (bind2 : bind1 : env') t - split2 sub tbind1 tbind2 body = - let (ptrs1', bs1') = split @env' tbind1 - bs1 = fst (weakenBindings weakenExpr WSink bs1') - (ptrs2, bs2) = split @(bind1 : env') tbind2 - in letBinds bs1 $ - letBinds (fst (weakenBindings weakenExpr (sinkWithBindings @(bind2 : bind1 : env') bs1) bs2)) $ - splitLets' (\cases _ IZ w -> subPointers ptrs2 (w .> wCopies (bindingsBinds bs2) (wSinks @(bind2 : bind1 : env') (bindingsBinds bs1))) - _ (IS IZ) w -> subPointers ptrs1' (w .> wSinks (bindingsBinds bs2) .> wCopies (bindingsBinds bs1) (WSink @bind2 @(bind1 : env'))) - t (IS (IS i)) w -> sub t i (WPop @bind1 (WPop @bind2 (wPops (bindingsBinds bs1) (wPops (bindingsBinds bs2) w))))) - body - -type family Split t where - Split (TPair a b) = SplitRec (TPair a b) - Split _ = '[] - -type family SplitRec t where - SplitRec TNil = '[] - SplitRec (TPair a b) = Append (SplitRec b) (SplitRec a) - SplitRec t = '[t] - -data Pointers env t where - Point :: STy t -> Idx env t -> Pointers env t - PNil :: Pointers env TNil - PPair :: Pointers env a -> Pointers env b -> Pointers env (TPair a b) - PWeak :: env' :> env -> Pointers env' t -> Pointers env t - -subPointers :: Pointers env t -> env :> env' -> Ex env' t -subPointers (Point t i) w = EVar ext t (w @> i) -subPointers PNil _ = ENil ext -subPointers (PPair a b) w = EPair ext (subPointers a w) (subPointers b w) -subPointers (PWeak w' p) w = subPointers p (w .> w') - -split :: forall env t. STy t - -> (Pointers (Append (Split t) (t : env)) t, Bindings Ex (t : env) (Split t)) -split typ = case typ of - STPair{} -> splitRec (EVar ext typ IZ) typ - STNil -> other - STEither{} -> other - STLEither{} -> other - STMaybe{} -> other - STArr{} -> other - STScal{} -> other - STAccum{} -> other - where - other :: (Pointers (t : env) t, Bindings Ex (t : env) '[]) - other = (Point typ IZ, BTop) - -splitRec :: forall env t. Ex env t -> STy t - -> (Pointers (Append (SplitRec t) env) t, Bindings Ex env (SplitRec t)) -splitRec rhs typ = case typ of - STNil -> (PNil, BTop) - STPair (a :: STy a) (b :: STy b) - | Refl <- lemAppendAssoc @(SplitRec b) @(SplitRec a) @env -> - let (p1, bs1) = splitRec (EFst ext rhs) a - (p2, bs2) = splitRec (ESnd ext (sinkWithBindings bs1 `weakenExpr` rhs)) b - in (PPair (PWeak (sinkWithBindings bs2) p1) p2, bconcat bs1 bs2) - STEither{} -> other - STLEither{} -> other - STMaybe{} -> other - STArr{} -> other - STScal{} -> other - STAccum{} -> other - where - other :: (Pointers (t : env) t, Bindings Ex env '[t]) - other = (Point typ IZ, BPush BTop (typ, rhs)) diff --git a/src/AST/Types.hs b/src/AST/Types.hs deleted file mode 100644 index 42bfb92..0000000 --- a/src/AST/Types.hs +++ /dev/null @@ -1,205 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE StandaloneKindSignatures #-} -{-# LANGUAGE TypeData #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -module AST.Types where - -import Data.Int (Int32, Int64) -import Data.GADT.Compare -import Data.GADT.Show -import Data.Kind (Type) -import Data.Type.Equality - -import Data - - -type data Ty - = TNil - | TPair Ty Ty - | TEither Ty Ty - | TLEither Ty Ty - | TMaybe Ty - | TArr Nat Ty -- ^ rank, element type - | TScal ScalTy - | TAccum Ty -- ^ contained type must be a monoid type - -type data ScalTy = TI32 | TI64 | TF32 | TF64 | TBool - -type STy :: Ty -> Type -data STy t where - STNil :: STy TNil - STPair :: STy a -> STy b -> STy (TPair a b) - STEither :: STy a -> STy b -> STy (TEither a b) - STLEither :: STy a -> STy b -> STy (TLEither a b) - STMaybe :: STy a -> STy (TMaybe a) - STArr :: SNat n -> STy t -> STy (TArr n t) - STScal :: SScalTy t -> STy (TScal t) - STAccum :: SMTy t -> STy (TAccum t) -deriving instance Show (STy t) - -instance GCompare STy where - gcompare = \cases - STNil STNil -> GEQ - STNil _ -> GLT ; _ STNil -> GGT - (STPair a b) (STPair a' b') -> gorderingLift2 (gcompare a a') (gcompare b b') - STPair{} _ -> GLT ; _ STPair{} -> GGT - (STEither a b) (STEither a' b') -> gorderingLift2 (gcompare a a') (gcompare b b') - STEither{} _ -> GLT ; _ STEither{} -> GGT - (STLEither a b) (STLEither a' b') -> gorderingLift2 (gcompare a a') (gcompare b b') - STLEither{} _ -> GLT ; _ STLEither{} -> GGT - (STMaybe a) (STMaybe a') -> gorderingLift1 (gcompare a a') - STMaybe{} _ -> GLT ; _ STMaybe{} -> GGT - (STArr n t) (STArr n' t') -> gorderingLift2 (gcompare n n') (gcompare t t') - STArr{} _ -> GLT ; _ STArr{} -> GGT - (STScal t) (STScal t') -> gorderingLift1 (gcompare t t') - STScal{} _ -> GLT ; _ STScal{} -> GGT - (STAccum t) (STAccum t') -> gorderingLift1 (gcompare t t') - -- STAccum{} _ -> GLT ; _ STAccum{} -> GGT - -instance TestEquality STy where testEquality = geq -instance GEq STy where geq = defaultGeq -instance GShow STy where gshowsPrec = defaultGshowsPrec - --- | Monoid types -type SMTy :: Ty -> Type -data SMTy t where - SMTNil :: SMTy TNil - SMTPair :: SMTy a -> SMTy b -> SMTy (TPair a b) - SMTLEither :: SMTy a -> SMTy b -> SMTy (TLEither a b) - SMTMaybe :: SMTy a -> SMTy (TMaybe a) - SMTArr :: SNat n -> SMTy t -> SMTy (TArr n t) - SMTScal :: ScalIsNumeric t ~ True => SScalTy t -> SMTy (TScal t) -deriving instance Show (SMTy t) - -instance GCompare SMTy where - gcompare = \cases - SMTNil SMTNil -> GEQ - SMTNil _ -> GLT ; _ SMTNil -> GGT - (SMTPair a b) (SMTPair a' b') -> gorderingLift2 (gcompare a a') (gcompare b b') - SMTPair{} _ -> GLT ; _ SMTPair{} -> GGT - (SMTLEither a b) (SMTLEither a' b') -> gorderingLift2 (gcompare a a') (gcompare b b') - SMTLEither{} _ -> GLT ; _ SMTLEither{} -> GGT - (SMTMaybe a) (SMTMaybe a') -> gorderingLift1 (gcompare a a') - SMTMaybe{} _ -> GLT ; _ SMTMaybe{} -> GGT - (SMTArr n t) (SMTArr n' t') -> gorderingLift2 (gcompare n n') (gcompare t t') - SMTArr{} _ -> GLT ; _ SMTArr{} -> GGT - (SMTScal t) (SMTScal t') -> gorderingLift1 (gcompare t t') - -- SMTScal{} _ -> GLT ; _ SMTScal{} -> GGT - -instance TestEquality SMTy where testEquality = geq -instance GEq SMTy where geq = defaultGeq -instance GShow SMTy where gshowsPrec = defaultGshowsPrec - -fromSMTy :: SMTy t -> STy t -fromSMTy = \case - SMTNil -> STNil - SMTPair t1 t2 -> STPair (fromSMTy t1) (fromSMTy t2) - SMTLEither t1 t2 -> STLEither (fromSMTy t1) (fromSMTy t2) - SMTMaybe t -> STMaybe (fromSMTy t) - SMTArr n t -> STArr n (fromSMTy t) - SMTScal sty -> STScal sty - -data SScalTy t where - STI32 :: SScalTy TI32 - STI64 :: SScalTy TI64 - STF32 :: SScalTy TF32 - STF64 :: SScalTy TF64 - STBool :: SScalTy TBool -deriving instance Show (SScalTy t) - -instance GCompare SScalTy where - gcompare = \cases - STI32 STI32 -> GEQ - STI32 _ -> GLT ; _ STI32 -> GGT - STI64 STI64 -> GEQ - STI64 _ -> GLT ; _ STI64 -> GGT - STF32 STF32 -> GEQ - STF32 _ -> GLT ; _ STF32 -> GGT - STF64 STF64 -> GEQ - STF64 _ -> GLT ; _ STF64 -> GGT - STBool STBool -> GEQ - -- STBool _ -> GLT ; _ STBool -> GGT - -instance TestEquality SScalTy where testEquality = geq -instance GEq SScalTy where geq = defaultGeq -instance GShow SScalTy where gshowsPrec = defaultGshowsPrec - -scalRepIsShow :: SScalTy t -> Dict (Show (ScalRep t)) -scalRepIsShow STI32 = Dict -scalRepIsShow STI64 = Dict -scalRepIsShow STF32 = Dict -scalRepIsShow STF64 = Dict -scalRepIsShow STBool = Dict - -type TIx = TScal TI64 - -tIx :: STy TIx -tIx = STScal STI64 - -type family ScalRep t where - ScalRep TI32 = Int32 - ScalRep TI64 = Int64 - ScalRep TF32 = Float - ScalRep TF64 = Double - ScalRep TBool = Bool - -type family ScalIsNumeric t where - ScalIsNumeric TI32 = True - ScalIsNumeric TI64 = True - ScalIsNumeric TF32 = True - ScalIsNumeric TF64 = True - ScalIsNumeric TBool = False - -type family ScalIsFloating t where - ScalIsFloating TI32 = False - ScalIsFloating TI64 = False - ScalIsFloating TF32 = True - ScalIsFloating TF64 = True - ScalIsFloating TBool = False - -type family ScalIsIntegral t where - ScalIsIntegral TI32 = True - ScalIsIntegral TI64 = True - ScalIsIntegral TF32 = False - ScalIsIntegral TF64 = False - ScalIsIntegral TBool = False - --- | Returns true for arrays /and/ accumulators. -hasArrays :: STy t' -> Bool -hasArrays STNil = False -hasArrays (STPair a b) = hasArrays a || hasArrays b -hasArrays (STEither a b) = hasArrays a || hasArrays b -hasArrays (STLEither a b) = hasArrays a || hasArrays b -hasArrays (STMaybe t) = hasArrays t -hasArrays STArr{} = True -hasArrays STScal{} = False -hasArrays STAccum{} = True - -type family Tup env where - Tup '[] = TNil - Tup (t : ts) = TPair (Tup ts) t - -mkTup :: f TNil -> (forall a b. f a -> f b -> f (TPair a b)) - -> SList f list -> f (Tup list) -mkTup nil _ SNil = nil -mkTup nil pair (e `SCons` es) = pair (mkTup nil pair es) e - -tTup :: SList STy env -> STy (Tup env) -tTup = mkTup STNil STPair - -unTup :: (forall a b. c (TPair a b) -> (c a, c b)) - -> SList f list -> c (Tup list) -> SList c list -unTup _ SNil _ = SNil -unTup unpack (_ `SCons` list) tup = - let (xs, x) = unpack tup - in x `SCons` unTup unpack list xs - -type family InvTup core env where - InvTup core '[] = core - InvTup core (t : ts) = InvTup (TPair core t) ts diff --git a/src/AST/UnMonoid.hs b/src/AST/UnMonoid.hs deleted file mode 100644 index ef01bf8..0000000 --- a/src/AST/UnMonoid.hs +++ /dev/null @@ -1,249 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE TypeOperators #-} -module AST.UnMonoid (unMonoid, zero, plus, acPrjCompose) where - -import AST -import AST.Sparse.Types -import Data - - --- | Remove 'EZero', 'EDeepZero', 'EPlus' and 'EOneHot' from the program by --- expanding them into their concrete implementations. Also ensure that --- 'EAccum' has a dense sparsity. -unMonoid :: Ex env t -> Ex env t -unMonoid = \case - EZero _ t e -> zero t e - EDeepZero _ t e -> deepZero t e - EPlus _ t a b -> plus t (unMonoid a) (unMonoid b) - EOneHot _ t p a b -> onehot t p (unMonoid a) (unMonoid b) - - EVar _ t i -> EVar ext t i - ELet _ rhs body -> ELet ext (unMonoid rhs) (unMonoid body) - EPair _ a b -> EPair ext (unMonoid a) (unMonoid b) - EFst _ e -> EFst ext (unMonoid e) - ESnd _ e -> ESnd ext (unMonoid e) - ENil _ -> ENil ext - EInl _ t e -> EInl ext t (unMonoid e) - EInr _ t e -> EInr ext t (unMonoid e) - ECase _ e a b -> ECase ext (unMonoid e) (unMonoid a) (unMonoid b) - ENothing _ t -> ENothing ext t - EJust _ e -> EJust ext (unMonoid e) - EMaybe _ a b e -> EMaybe ext (unMonoid a) (unMonoid b) (unMonoid e) - ELNil _ t1 t2 -> ELNil ext t1 t2 - ELInl _ t e -> ELInl ext t (unMonoid e) - ELInr _ t e -> ELInr ext t (unMonoid e) - ELCase _ e a b c -> ELCase ext (unMonoid e) (unMonoid a) (unMonoid b) (unMonoid c) - EConstArr _ n t x -> EConstArr ext n t x - EBuild _ n a b -> EBuild ext n (unMonoid a) (unMonoid b) - EFold1Inner _ cm a b c -> EFold1Inner ext cm (unMonoid a) (unMonoid b) (unMonoid c) - ESum1Inner _ e -> ESum1Inner ext (unMonoid e) - EUnit _ e -> EUnit ext (unMonoid e) - EReplicate1Inner _ a b -> EReplicate1Inner ext (unMonoid a) (unMonoid b) - EMaximum1Inner _ e -> EMaximum1Inner ext (unMonoid e) - EMinimum1Inner _ e -> EMinimum1Inner ext (unMonoid e) - EConst _ t x -> EConst ext t x - EIdx0 _ e -> EIdx0 ext (unMonoid e) - EIdx1 _ a b -> EIdx1 ext (unMonoid a) (unMonoid b) - EIdx _ a b -> EIdx ext (unMonoid a) (unMonoid b) - EShape _ e -> EShape ext (unMonoid e) - EOp _ op e -> EOp ext op (unMonoid e) - ECustom _ t1 t2 t3 a b c e1 e2 -> ECustom ext t1 t2 t3 (unMonoid a) (unMonoid b) (unMonoid c) (unMonoid e1) (unMonoid e2) - ERecompute _ e -> ERecompute ext (unMonoid e) - EWith _ t a b -> EWith ext t (unMonoid a) (unMonoid b) - EAccum _ t p eidx sp eval eacc -> - accumulateSparse (acPrjTy p t) sp eval $ \w prj2 idx2 val2 -> - acPrjCompose SAID p (weakenExpr w eidx) prj2 idx2 $ \prj' idx' -> - EAccum ext t prj' (unMonoid idx') (spDense (acPrjTy prj' t)) (unMonoid val2) (weakenExpr w (unMonoid eacc)) - EError _ t s -> EError ext t s - -zero :: SMTy t -> Ex env (ZeroInfo t) -> Ex env t -zero SMTNil e = elet e $ ENil ext -zero (SMTPair t1 t2) e = - ELet ext e $ EPair ext (zero t1 (EFst ext (EVar ext (typeOf e) IZ))) - (zero t2 (ESnd ext (EVar ext (typeOf e) IZ))) -zero (SMTLEither t1 t2) _ = ELNil ext (fromSMTy t1) (fromSMTy t2) -zero (SMTMaybe t) _ = ENothing ext (fromSMTy t) -zero (SMTArr _ t) e = emap (zero t (EVar ext (tZeroInfo t) IZ)) e -zero (SMTScal t) _ = case t of - STI32 -> EConst ext STI32 0 - STI64 -> EConst ext STI64 0 - STF32 -> EConst ext STF32 0.0 - STF64 -> EConst ext STF64 0.0 - -deepZero :: SMTy t -> Ex env (DeepZeroInfo t) -> Ex env t -deepZero SMTNil e = elet e $ ENil ext -deepZero (SMTPair t1 t2) e = - ELet ext e $ EPair ext (deepZero t1 (EFst ext (EVar ext (typeOf e) IZ))) - (deepZero t2 (ESnd ext (EVar ext (typeOf e) IZ))) -deepZero (SMTLEither t1 t2) e = - elcase e - (ELNil ext (fromSMTy t1) (fromSMTy t2)) - (ELInl ext (fromSMTy t2) (deepZero t1 (evar IZ))) - (ELInr ext (fromSMTy t1) (deepZero t2 (evar IZ))) -deepZero (SMTMaybe t) e = - emaybe e - (ENothing ext (fromSMTy t)) - (EJust ext (deepZero t (evar IZ))) -deepZero (SMTArr _ t) e = emap (deepZero t (evar IZ)) e -deepZero (SMTScal t) _ = case t of - STI32 -> EConst ext STI32 0 - STI64 -> EConst ext STI64 0 - STF32 -> EConst ext STF32 0.0 - STF64 -> EConst ext STF64 0.0 - -plus :: SMTy t -> Ex env t -> Ex env t -> Ex env t --- don't destroy the effects! -plus SMTNil a b = elet a $ elet (weakenExpr WSink b) $ ENil ext -plus (SMTPair t1 t2) a b = - let t = STPair (fromSMTy t1) (fromSMTy t2) - in ELet ext a $ - ELet ext (weakenExpr WSink b) $ - EPair ext (plus t1 (EFst ext (EVar ext t (IS IZ))) - (EFst ext (EVar ext t IZ))) - (plus t2 (ESnd ext (EVar ext t (IS IZ))) - (ESnd ext (EVar ext t IZ))) -plus (SMTLEither t1 t2) a b = - let t = STLEither (fromSMTy t1) (fromSMTy t2) - in ELet ext a $ - ELet ext (weakenExpr WSink b) $ - ELCase ext (EVar ext t (IS IZ)) - (EVar ext t IZ) - (ELCase ext (EVar ext t (IS IZ)) - (EVar ext t (IS (IS IZ))) - (ELInl ext (fromSMTy t2) (plus t1 (EVar ext (fromSMTy t1) (IS IZ)) (EVar ext (fromSMTy t1) IZ))) - (EError ext t "plus l+r")) - (ELCase ext (EVar ext t (IS IZ)) - (EVar ext t (IS (IS IZ))) - (EError ext t "plus r+l") - (ELInr ext (fromSMTy t1) (plus t2 (EVar ext (fromSMTy t2) (IS IZ)) (EVar ext (fromSMTy t2) IZ)))) -plus (SMTMaybe t) a b = - ELet ext b $ - EMaybe ext - (EVar ext (STMaybe (fromSMTy t)) IZ) - (EJust ext - (EMaybe ext - (EVar ext (fromSMTy t) IZ) - (plus t (EVar ext (fromSMTy t) (IS IZ)) (EVar ext (fromSMTy t) IZ)) - (EVar ext (STMaybe (fromSMTy t)) (IS IZ)))) - (weakenExpr WSink a) -plus (SMTArr _ t) a b = - ezipWith (plus t (EVar ext (fromSMTy t) (IS IZ)) (EVar ext (fromSMTy t) IZ)) - a b -plus (SMTScal t) a b = EOp ext (OAdd t) (EPair ext a b) - -onehot :: SMTy t -> SAcPrj p t a -> Ex env (AcIdxS p t) -> Ex env a -> Ex env t -onehot typ topprj idx arg = case (typ, topprj) of - (_, SAPHere) -> - ELet ext arg $ - EVar ext (fromSMTy typ) IZ - - (SMTPair t1 t2, SAPFst prj) -> - ELet ext idx $ - let tidx = typeOf idx in - ELet ext (onehot t1 prj (EFst ext (EVar ext (typeOf idx) IZ)) (weakenExpr WSink arg)) $ - let toh = fromSMTy t1 in - EPair ext (EVar ext toh IZ) - (zero t2 (ESnd ext (EVar ext tidx (IS IZ)))) - - (SMTPair t1 t2, SAPSnd prj) -> - ELet ext idx $ - let tidx = typeOf idx in - ELet ext (onehot t2 prj (ESnd ext (EVar ext (typeOf idx) IZ)) (weakenExpr WSink arg)) $ - let toh = fromSMTy t2 in - EPair ext (zero t1 (EFst ext (EVar ext tidx (IS IZ)))) - (EVar ext toh IZ) - - (SMTLEither t1 t2, SAPLeft prj) -> - ELInl ext (fromSMTy t2) (onehot t1 prj idx arg) - (SMTLEither t1 t2, SAPRight prj) -> - ELInr ext (fromSMTy t1) (onehot t2 prj idx arg) - - (SMTMaybe t1, SAPJust prj) -> - EJust ext (onehot t1 prj idx arg) - - (SMTArr n t1, SAPArrIdx prj) -> - let tidx = tTup (sreplicate n tIx) - in ELet ext idx $ - EBuild ext n (EShape ext (ESnd ext (EFst ext (EVar ext (typeOf idx) IZ)))) $ - eif (eidxEq n (EVar ext tidx IZ) (EFst ext (EFst ext (EVar ext (typeOf idx) (IS IZ))))) - (onehot t1 prj (ESnd ext (EVar ext (typeOf idx) (IS IZ))) (weakenExpr (WSink .> WSink) arg)) - (ELet ext (EIdx ext (ESnd ext (EFst ext (EVar ext (typeOf idx) (IS IZ)))) (EVar ext tidx IZ)) $ - zero t1 (EVar ext (tZeroInfo t1) IZ)) - -accumulateSparse - :: SMTy t -> Sparse t t' -> Ex env t' - -> (forall p b env'. env :> env' -> SAcPrj p t b -> Ex env' (AcIdxD p t) -> Ex env' b -> Ex env' TNil) - -> Ex env TNil -accumulateSparse topty topsp arg accum = case (topty, topsp) of - (_, s) | Just Refl <- isDense topty s -> - accum WId SAPHere (ENil ext) arg - (SMTScal _, SpScal) -> - accum WId SAPHere (ENil ext) arg -- should be handled by isDense already, but meh - (_, SpSparse s) -> - emaybe arg - (ENil ext) - (accumulateSparse topty s (evar IZ) (\w -> accum (WPop w))) - (_, SpAbsent) -> - ENil ext - (SMTPair t1 t2, SpPair s1 s2) -> - eunPair arg $ \w1 e1 e2 -> - elet (accumulateSparse t1 s1 e1 (\w prj -> accum (w .> w1) (SAPFst prj))) $ - accumulateSparse t2 s2 (weakenExpr WSink e2) (\w prj -> accum (w .> WSink .> w1) (SAPSnd prj)) - (SMTLEither t1 t2, SpLEither s1 s2) -> - elcase arg - (ENil ext) - (accumulateSparse t1 s1 (evar IZ) (\w prj -> accum (WPop w) (SAPLeft prj))) - (accumulateSparse t2 s2 (evar IZ) (\w prj -> accum (WPop w) (SAPRight prj))) - (SMTMaybe t, SpMaybe s) -> - emaybe arg - (ENil ext) - (accumulateSparse t s (evar IZ) (\w prj -> accum (WPop w) (SAPJust prj))) - (SMTArr n t, SpArr s) -> - let tn = tTup (sreplicate n tIx) in - elet arg $ - elet (EBuild ext n (EShape ext (evar IZ)) $ - accumulateSparse t s - (EIdx ext (evar (IS IZ)) (EVar ext tn IZ)) - (\w prj idx val -> accum (WPop (WPop w)) (SAPArrIdx prj) (EPair ext (EVar ext tn (w @> IZ)) idx) val)) $ - ENil ext - -acPrjCompose - :: SAIDense dense - -> SAcPrj p1 a b -> Ex env (AcIdx dense p1 a) - -> SAcPrj p2 b c -> Ex env (AcIdx dense p2 b) - -> (forall p'. SAcPrj p' a c -> Ex env (AcIdx dense p' a) -> r) -> r -acPrjCompose _ SAPHere _ p2 idx2 k = k p2 idx2 -acPrjCompose SAID (SAPFst p1) idx1 p2 idx2 k = - acPrjCompose SAID p1 idx1 p2 idx2 $ \p' idx' -> - k (SAPFst p') idx' -acPrjCompose SAID (SAPSnd p1) idx1 p2 idx2 k = - acPrjCompose SAID p1 idx1 p2 idx2 $ \p' idx' -> - k (SAPSnd p') idx' -acPrjCompose SAIS (SAPFst p1) idx1 p2 idx2 k - | Dict <- styKnown (typeOf idx1) = - acPrjCompose SAIS p1 (efst (evar IZ)) p2 (weakenExpr WSink idx2) $ \p' idx' -> - k (SAPFst p') (elet idx1 $ EPair ext idx' (esnd (evar IZ))) -acPrjCompose SAIS (SAPSnd p1) idx1 p2 idx2 k - | Dict <- styKnown (typeOf idx1) = - acPrjCompose SAIS p1 (esnd (evar IZ)) p2 (weakenExpr WSink idx2) $ \p' idx' -> - k (SAPSnd p') (elet idx1 $ EPair ext (efst (evar IZ)) idx') -acPrjCompose d (SAPLeft p1) idx1 p2 idx2 k = - acPrjCompose d p1 idx1 p2 idx2 $ \p' idx' -> - k (SAPLeft p') idx' -acPrjCompose d (SAPRight p1) idx1 p2 idx2 k = - acPrjCompose d p1 idx1 p2 idx2 $ \p' idx' -> - k (SAPRight p') idx' -acPrjCompose d (SAPJust p1) idx1 p2 idx2 k = - acPrjCompose d p1 idx1 p2 idx2 $ \p' idx' -> - k (SAPJust p') idx' -acPrjCompose SAID (SAPArrIdx p1) idx1 p2 idx2 k - | Dict <- styKnown (typeOf idx1) = - acPrjCompose SAID p1 (esnd (evar IZ)) p2 (weakenExpr WSink idx2) $ \p' idx' -> - k (SAPArrIdx p') (elet idx1 $ EPair ext (EFst ext (EVar ext (typeOf idx1) IZ)) idx') -acPrjCompose SAIS (SAPArrIdx p1) idx1 p2 idx2 k - | Dict <- styKnown (typeOf idx1) = - acPrjCompose SAIS p1 (esnd (evar IZ)) p2 (weakenExpr WSink idx2) $ \p' idx' -> - k (SAPArrIdx p') (elet idx1 $ EPair ext (EFst ext (EVar ext (typeOf idx1) IZ)) idx') diff --git a/src/AST/Weaken.hs b/src/AST/Weaken.hs deleted file mode 100644 index d882e28..0000000 --- a/src/AST/Weaken.hs +++ /dev/null @@ -1,132 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE EmptyCase #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE StandaloneKindSignatures #-} -{-# LANGUAGE TypeAbstractions #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} - -{-# LANGUAGE PartialTypeSignatures #-} -{-# OPTIONS -Wno-partial-type-signatures #-} - --- The reason why this is a separate module with "little" in it: -{-# LANGUAGE AllowAmbiguousTypes #-} - -module AST.Weaken (module AST.Weaken, Append) where - -import Data.Bifunctor (first) -import Data.Functor.Const -import Data.Kind (Type) - -import Data -import Lemmas - - -type Idx :: [k] -> k -> Type -data Idx env t where - IZ :: Idx (t : env) t - IS :: Idx env t -> Idx (a : env) t -deriving instance Show (Idx env t) - -splitIdx :: forall env2 env1 t f. SList f env1 -> Idx (Append env1 env2) t -> Either (Idx env1 t) (Idx env2 t) -splitIdx SNil i = Right i -splitIdx (SCons _ _) IZ = Left IZ -splitIdx (SCons _ l) (IS i) = first IS (splitIdx l i) - -slistIdx :: SList f list -> Idx list t -> f t -slistIdx (SCons x _) IZ = x -slistIdx (SCons _ list) (IS i) = slistIdx list i -slistIdx SNil i = case i of {} - -idx2int :: Idx env t -> Int -idx2int IZ = 0 -idx2int (IS n) = 1 + idx2int n - -data env :> env' where - WId :: env :> env - WSink :: forall t env. env :> (t : env) - WCopy :: forall t env env'. env :> env' -> (t : env) :> (t : env') - WPop :: (t : env) :> env' -> env :> env' - WThen :: env1 :> env2 -> env2 :> env3 -> env1 :> env3 - WClosed :: '[] :> env - WIdx :: Idx env t -> (t : env) :> env - WPick :: forall t pre env env'. SList (Const ()) pre -> env :> env' - -> Append pre (t : env) :> t : Append pre env' - WSwap :: forall env as bs. SList (Const ()) as -> SList (Const ()) bs - -> Append as (Append bs env) :> Append bs (Append as env) - WStack :: forall env1 env2 as bs. SList (Const ()) as -> SList (Const ()) bs - -> as :> bs -> env1 :> env2 - -> Append as env1 :> Append bs env2 -deriving instance Show (env :> env') -infix 4 :> - -infixr 2 @> -(@>) :: env :> env' -> Idx env t -> Idx env' t -WId @> i = i -WSink @> i = IS i -WCopy _ @> IZ = IZ -WCopy w @> IS i = IS (w @> i) -WPop w @> i = w @> IS i -WThen w1 w2 @> i = w2 @> w1 @> i -WClosed @> i = case i of {} -WIdx j @> IZ = j -WIdx _ @> IS i = i -WPick SNil w @> i = WCopy w @> i -WPick (_ `SCons` _) _ @> IZ = IS IZ -WPick @t (_ `SCons` pre) w @> IS i = WCopy WSink .> WPick @t pre w @> i -WSwap @env (as :: SList _ as) (bs :: SList _ bs) @> i = - case splitIdx @(Append bs env) as i of - Left i' -> indexSinks bs (indexRaiseAbove @env as i') - Right i' -> case splitIdx @env bs i' of - Left j -> indexRaiseAbove @(Append as env) bs j - Right j -> indexSinks bs (indexSinks as j) -WStack @env1 @env2 as bs wlo whi @> i = - case splitIdx @env1 as i of - Left i' -> indexRaiseAbove @env2 bs (wlo @> i') - Right i' -> indexSinks bs (whi @> i') - -indexSinks :: SList f as -> Idx bs t -> Idx (Append as bs) t -indexSinks SNil j = j -indexSinks (_ `SCons` bs') j = IS (indexSinks bs' j) - -indexRaiseAbove :: forall env as t f. SList f as -> Idx as t -> Idx (Append as env) t -indexRaiseAbove = flip go - where - go :: forall as'. Idx as' t -> SList f as' -> Idx (Append as' env) t - go IZ (_ `SCons` _) = IZ - go (IS i) (_ `SCons` as) = IS (go i as) - -infixr 3 .> -(.>) :: env2 :> env3 -> env1 :> env2 -> env1 :> env3 -(.>) = flip WThen - -class KnownListSpine list where knownListSpine :: SList (Const ()) list -instance KnownListSpine '[] where knownListSpine = SNil -instance KnownListSpine list => KnownListSpine (t : list) where knownListSpine = SCons (Const ()) knownListSpine - -wSinks' :: forall list env. KnownListSpine list => env :> Append list env -wSinks' = wSinks (knownListSpine :: SList (Const ()) list) - -wSinks :: forall env bs f. SList f bs -> env :> Append bs env -wSinks SNil = WId -wSinks (SCons _ spine) = WSink .> wSinks spine - -wSinksAnd :: forall env env' bs f. SList f bs -> env :> env' -> env :> Append bs env' -wSinksAnd SNil w = w -wSinksAnd (SCons _ spine) w = WSink .> wSinksAnd spine w - -wCopies :: SList f bs -> env1 :> env2 -> Append bs env1 :> Append bs env2 -wCopies bs w = - let bs' = slistMap (\_ -> Const ()) bs - in WStack bs' bs' WId w - -wRaiseAbove :: SList f env1 -> SList g env -> env1 :> Append env1 env -wRaiseAbove SNil _ = WClosed -wRaiseAbove (SCons _ s) env = WCopy (wRaiseAbove s env) - -wPops :: SList f bs -> Append bs env1 :> env2 -> env1 :> env2 -wPops SNil w = w -wPops (_ `SCons` bs) w = wPops bs (WPop w) diff --git a/src/AST/Weaken/Auto.hs b/src/AST/Weaken/Auto.hs deleted file mode 100644 index c6efe37..0000000 --- a/src/AST/Weaken/Auto.hs +++ /dev/null @@ -1,176 +0,0 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE FunctionalDependencies #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE KindSignatures #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeAbstractions #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} - -{-# LANGUAGE AllowAmbiguousTypes #-} - -{-# LANGUAGE PartialTypeSignatures #-} -{-# OPTIONS_GHC -Wno-partial-type-signatures #-} -module AST.Weaken.Auto ( - autoWeak, - (&.), auto, auto1, - Layout(..), -) where - -import Data.Functor.Const -import GHC.OverloadedLabels -import GHC.TypeLits -import Unsafe.Coerce (unsafeCoerce) - -import AST.Weaken -import Data -import Lemmas - - -type family Lookup name list where - Lookup name ('(name, x) : _) = x - Lookup name (_ : list) = Lookup name list - Lookup name '[] = TypeError (Text "The name '" :<>: Text name :<>: Text "' does not appear in the list.") - - --- | The @withPre@ type parameter indicates whether there can be 'LPreW' --- occurrences within this layout. -data Layout (withPre :: Bool) (segments :: [(Symbol, [t])]) (env :: [t]) where - LSeg :: forall name segments withPre. SSymbol name -> Layout withPre segments (Lookup name segments) - -- | Pre-weaken with a weakening - LPreW :: forall name1 name2 segments. - SegmentName name1 -> SegmentName name2 - -> Lookup name1 segments :> Lookup name2 segments - -> Layout True segments (Lookup name1 segments) - (:++:) :: Layout withPre segments env1 -> Layout withPre segments env2 -> Layout withPre segments (Append env1 env2) -infixr :++: - -instance (KnownSymbol name, seg ~ Lookup name segments) => IsLabel name (Layout withPre segments seg) where - fromLabel = LSeg (symbolSing @name) - -newtype SegmentName name = SegmentName (SSymbol name) - deriving (Show) - -instance (KnownSymbol name, name ~ name') => IsLabel name (SegmentName name') where - fromLabel = SegmentName symbolSing - - -data SSegments (segments :: [(Symbol, [t])]) where - SSegNil :: SSegments '[] - SSegCons :: SSymbol name -> SList (Const ()) ts -> SSegments list -> SSegments ('(name, ts) : list) - -instance (KnownSymbol name, segs ~ '[ '(name, ts)]) => IsLabel name (SList f ts -> SSegments segs) where - fromLabel = \spine -> SSegCons symbolSing (slistMap (\_ -> Const ()) spine) SSegNil - -auto :: KnownListSpine list => SList (Const ()) list -auto = knownListSpine - -auto1 :: SList (Const ()) '[t] -auto1 = Const () `SCons` SNil - -infixr &. -(&.) :: SSegments segs1 -> SSegments segs2 -> SSegments (Append segs1 segs2) -(&.) = ssegmentsAppend - where - ssegmentsAppend :: SSegments a -> SSegments b -> SSegments (Append a b) - ssegmentsAppend SSegNil l2 = l2 - ssegmentsAppend (SSegCons name list l1) l2 = SSegCons name list (ssegmentsAppend l1 l2) - - --- | If the found segment is a TopSeg, returns Nothing. -segmentLookup :: forall segments name. SSegments segments -> SSymbol name -> SList (Const ()) (Lookup name segments) -segmentLookup = \segs name -> case go segs name of - Just ts -> ts - Nothing -> error $ "Segment not found: " ++ fromSSymbol name - where - go :: forall segs'. SSegments segs' -> SSymbol name -> Maybe (SList (Const ()) (Lookup name segs')) - go SSegNil _ = Nothing - go (SSegCons n@(SSymbol @n) (ts :: SList _ ts) (sseg :: SSegments rest)) name@SSymbol = - case sameSymbol n name of - Just Refl -> - case go sseg name of - Nothing -> Just ts - Just _ -> error $ "Duplicate segment with name " ++ fromSSymbol name - Nothing -> - case unsafeCoerce Refl :: (Lookup name ('(n, ts) : rest) :~: Lookup name rest) of - Refl -> go sseg name - -data LinLayout (withPre :: Bool) (segments :: [(Symbol, [t])]) (env :: [t]) where - LinEnd :: LinLayout withPre segments '[] - LinApp :: SSymbol name -> LinLayout withPre segments env - -> LinLayout withPre segments (Append (Lookup name segments) env) - LinAppPreW :: SSymbol name1 -> SSymbol name2 - -> Lookup name1 segments :> Lookup name2 segments - -> LinLayout True segments env - -> LinLayout True segments (Append (Lookup name1 segments) env) - -linLayoutAppend :: LinLayout withPre segments env1 -> LinLayout withPre segments env2 -> LinLayout withPre segments (Append env1 env2) -linLayoutAppend LinEnd lin = lin -linLayoutAppend (LinApp (name :: SSymbol name) (lin1 :: LinLayout _ segments env1')) (lin2 :: LinLayout _ _ env2) - | Refl <- lemAppendAssoc @(Lookup name segments) @env1' @env2 - = LinApp name (linLayoutAppend lin1 lin2) -linLayoutAppend (LinAppPreW (name1 :: SSymbol name1) name2 w (lin1 :: LinLayout _ segments env1')) (lin2 :: LinLayout _ _ env2) - | Refl <- lemAppendAssoc @(Lookup name1 segments) @env1' @env2 - = LinAppPreW name1 name2 w (linLayoutAppend lin1 lin2) - -lineariseLayout :: Layout withPre segments env -> LinLayout withPre segments env -lineariseLayout (LSeg name :: Layout _ _ seg) - | Refl <- lemAppendNil @seg - = LinApp name LinEnd -lineariseLayout (ly1 :++: ly2) = lineariseLayout ly1 `linLayoutAppend` lineariseLayout ly2 -lineariseLayout (LPreW (SegmentName name1) (SegmentName name2) w :: Layout _ _ seg) - | Refl <- lemAppendNil @seg - = LinAppPreW name1 name2 w LinEnd - -preWeaken :: SSegments segments -> LinLayout True segments env - -> (forall env'. env :> env' -> LinLayout False segments env' -> r) -> r -preWeaken _ LinEnd k = k WId LinEnd -preWeaken segs (LinApp name lin) k = - preWeaken segs lin $ \w lin' -> - k (wCopies (segmentLookup segs name) w) (LinApp name lin') -preWeaken segs (LinAppPreW name1 name2 weak lin) k = - preWeaken segs lin $ \w lin' -> - k (WStack (segmentLookup segs name1) (segmentLookup segs name2) weak w) (LinApp name2 lin') - -pullDown :: SSegments segments -> SSymbol name -> LinLayout False segments env - -> r -- Name was not found in source - -> (forall env'. LinLayout False segments env' -> env :> Append (Lookup name segments) env' -> r) - -> r -pullDown segs name@SSymbol linlayout kNotFound k = - case linlayout of - LinEnd -> kNotFound - LinApp n'@SSymbol lin - | Just Refl <- sameSymbol name n' -> k lin WId - | otherwise -> - pullDown segs name lin kNotFound $ \(lin' :: LinLayout _ _ env') w -> - k (LinApp n' lin') (WSwap @env' (segmentLookup segs n') (segmentLookup segs name) - .> wCopies (segmentLookup segs n') w) - -sortLinLayouts :: forall segments env1 env2. - SSegments segments - -> LinLayout False segments env1 -> LinLayout False segments env2 -> env1 :> env2 -sortLinLayouts _ LinEnd LinEnd = WId -sortLinLayouts segs lin1@(LinApp name1@SSymbol tail1) (LinApp name2@SSymbol tail2) - | Just Refl <- sameSymbol name1 name2 = wCopies (segmentLookup segs name1) (sortLinLayouts segs tail1 tail2) - | otherwise = - pullDown segs name2 lin1 - (wSinks (segmentLookup segs name2) .> sortLinLayouts segs lin1 tail2) - (\tail1' w -> - -- We've pulled down name2 in lin1 so that it's at the head; the - -- resulting modified tail is tail1'. Thus now we have (name2 : tail1') - -- vs (name2 : tail2). Thus we continue sorting tail1' vs tail2, and - -- wCopies the name2 on top of that. - wCopies (segmentLookup segs name2) (sortLinLayouts segs tail1' tail2) .> w) -sortLinLayouts _ LinEnd LinApp{} = WClosed -sortLinLayouts _ LinApp{} LinEnd = error "Segments in source that do not occur in target" - -autoWeak :: forall segments env1 env2. - SSegments segments -> Layout True segments env1 -> Layout False segments env2 -> env1 :> env2 -autoWeak segs ly1 ly2 = - preWeaken segs (lineariseLayout ly1) $ \wPreweak lin1 -> - sortLinLayouts segs lin1 (lineariseLayout ly2) .> wPreweak |