Skip to content

Commit

Permalink
Removing colors and IRRep
Browse files Browse the repository at this point in the history
  • Loading branch information
dougalm committed May 17, 2024
1 parent d0334fc commit 0e6e527
Show file tree
Hide file tree
Showing 25 changed files with 1,738 additions and 2,327 deletions.
1 change: 0 additions & 1 deletion dex.cabal
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,6 @@ library
, IncState
, Inference
-- , Inline
, IRVariants
-- , JAX.Concrete
-- , JAX.Rename
-- , JAX.ToSimp
Expand Down
2 changes: 1 addition & 1 deletion src/lib/AbstractSyntax.hs
Original file line number Diff line number Diff line change
Expand Up @@ -255,7 +255,7 @@ patOptAnn (WithSrcs _ _ (CBin Colon lhs typeAnn)) = (,) <$> pat lhs <*> (Just <$
patOptAnn (WithSrcs _ _ (CParens [g])) = patOptAnn g
patOptAnn g = (,Nothing) <$> pat g

uBinder :: GroupW -> SyntaxM (UBinder c VoidS VoidS)
uBinder :: GroupW -> SyntaxM (UBinder VoidS VoidS)
uBinder (WithSrcs sid _ b) = case b of
CLeaf (CIdentifier name) -> return $ fromSourceNameW $ WithSrc sid name
CLeaf CHole -> return $ WithSrcB sid UIgnore
Expand Down
554 changes: 275 additions & 279 deletions src/lib/Builder.hs

Large diffs are not rendered by default.

207 changes: 104 additions & 103 deletions src/lib/CheapReduction.hs

Large diffs are not rendered by default.

1 change: 0 additions & 1 deletion src/lib/ConcreteSyntax.hs
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@ import Text.Megaparsec.Char hiding (space, eol)

import Err
import Lexing
import Types.Core
import Types.Source
import Types.Primitives
import Util
Expand Down
48 changes: 22 additions & 26 deletions src/lib/Core.hs
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,6 @@ import qualified Data.Map.Strict as M

import Name
import Err
import IRVariants

import Types.Core
import Types.Top
Expand Down Expand Up @@ -189,20 +188,20 @@ class BindsNames b => BindsEnv (b::B) where
=> Distinct l => b n l -> EnvFrag n l
toEnvFrag b = toEnvFrag $ fromB b

instance (Color c, SinkableE ann, ToBinding ann c) => BindsEnv (BinderP c ann) where
instance (SinkableE ann, ToBinding ann) => BindsEnv (BinderP ann) where
toEnvFrag (b:>ann) = EnvFrag (RecSubstFrag (b @> toBinding ann'))
where ann' = withExtEvidence b $ sink ann

instance (IRRep r, SinkableE ann, ToBinding ann (AtomNameC r)) => BindsEnv (NonDepNest r ann) where
instance (SinkableE ann, ToBinding ann) => BindsEnv (NonDepNest ann) where
toEnvFrag (NonDepNest topBs topAnns) = toEnvFrag $ zipNest topBs topAnns
where
zipNest :: Distinct l => Nest (AtomNameBinder r) n l -> [ann n] -> Nest (BinderP (AtomNameC r) ann) n l
zipNest :: Distinct l => Nest AtomNameBinder n l -> [ann n] -> Nest (BinderP ann) n l
zipNest Empty [] = Empty
zipNest (Nest b bs) (a:anns) = withExtEvidence b $ withSubscopeDistinct bs $
Nest (b:>a) $ zipNest bs $ sinkList anns
zipNest _ _ = error "Mismatched lengths in NonDepNest"

instance IRRep r => BindsEnv (Decl r) where
instance BindsEnv Decl where
toEnvFrag (Let b binding) = toEnvFrag $ b :> binding
{-# INLINE toEnvFrag #-}

Expand Down Expand Up @@ -265,12 +264,12 @@ instance (BindsEnv b1, BindsEnv b2)
instance BindsEnv UnitB where
toEnvFrag UnitB = emptyOutFrag

instance IRRep r => ExtOutMap Env (Nest (Decl r)) where
instance ExtOutMap Env (Nest Decl) where
extendOutMap bindings emissions =
bindings `extendOutMap` toEnvFrag emissions
{-# INLINE extendOutMap #-}

instance IRRep r => ExtOutMap Env (RNest (Decl r)) where
instance ExtOutMap Env (RNest Decl) where
extendOutMap bindings emissions =
bindings `extendOutMap` toEnvFrag emissions
{-# INLINE extendOutMap #-}
Expand All @@ -281,15 +280,15 @@ instance ExtOutMap Env UnitB where

-- === Monadic helpers ===

lookupEnv :: (Color c, EnvReader m) => Name c o -> m o (Binding c o)
lookupEnv :: EnvReader m => Name o -> m o (Binding o)
lookupEnv v = withEnv $ flip lookupEnvPure v . topEnv
{-# INLINE lookupEnv #-}

lookupAtomName :: (IRRep r, EnvReader m) => AtomName r n -> m n (AtomBinding r n)
lookupAtomName :: EnvReader m => AtomName n -> m n (AtomBinding n)
lookupAtomName name = bindingToAtomBinding <$> lookupEnv name
{-# INLINE lookupAtomName #-}

lookupCustomRules :: EnvReader m => AtomName CoreIR n -> m n (Maybe (AtomRules n))
lookupCustomRules :: EnvReader m => AtomName n -> m n (Maybe (AtomRules n))
lookupCustomRules name = liftM fromMaybeE $ withEnv $
toMaybeE . M.lookup name . customRulesMap . envCustomRules . topEnv
{-# INLINE lookupCustomRules #-}
Expand Down Expand Up @@ -317,7 +316,7 @@ lookupTyCon name = lookupEnv name >>= \case
TyConBinding Nothing _ -> error "TyCon not yet defined"
{-# INLINE lookupTyCon #-}

lookupDataCon :: EnvReader m => Name DataConNameC n -> m n (TyConName n, Int)
lookupDataCon :: EnvReader m => DataConName n -> m n (TyConName n, Int)
lookupDataCon v = do
~(DataConBinding defName idx) <- lookupEnv v
return (defName, idx)
Expand Down Expand Up @@ -356,19 +355,19 @@ refreshBinders b cont = refreshAbs (Abs b $ idSubstFrag b) cont
{-# INLINE refreshBinders #-}

withFreshBinder
:: (Color c, EnvExtender m, ToBinding binding c)
:: (EnvExtender m, ToBinding binding)
=> NameHint -> binding n
-> (forall l. DExt n l => BinderP c binding n l -> m l a)
-> (forall l. DExt n l => BinderP binding n l -> m l a)
-> m n a
withFreshBinder hint binding cont = do
Abs b v <- freshNameM hint
refreshAbs (Abs (b:>binding) v) \b' _ -> cont b'
{-# INLINE withFreshBinder #-}

withFreshBinders
:: (Color c, EnvExtender m, ToBinding binding c)
:: (EnvExtender m, ToBinding binding)
=> [binding n]
-> (forall l. DExt n l => Nest (BinderP c binding) n l -> [Name c l] -> m l a)
-> (forall l. DExt n l => Nest (BinderP binding) n l -> [Name l] -> m l a)
-> m n a
withFreshBinders [] cont = do
Distinct <- getDistinct
Expand All @@ -391,14 +390,14 @@ withFreshBinders (binding:rest) cont = do
-- present, in which case exactly maxDepth binders are packed into the nary
-- structure. Excess binders, if any, are still left in the unary structures.

liftLamExpr :: (IRRep r, EnvReader m)
=> TopLam r n
-> (forall l m2. EnvReader m2 => Expr r l -> m2 l (Expr r l))
-> m n (TopLam r n)
liftLamExpr :: EnvReader m
=> TopLam n
-> (forall l m2. EnvReader m2 => Expr l -> m2 l (Expr l))
-> m n (TopLam n)
liftLamExpr (TopLam d ty (LamExpr bs body)) f = liftM (TopLam d ty) $ liftEnvReaderM $
refreshAbs (Abs bs body) \bs' body' -> LamExpr bs' <$> f body'

fromNaryForExpr :: IRRep r => Int -> Expr r n -> Maybe (Int, LamExpr r n)
fromNaryForExpr :: Int -> Expr n -> Maybe (Int, LamExpr n)
fromNaryForExpr maxDepth | maxDepth <= 0 = error "expected non-negative number of args"
fromNaryForExpr maxDepth = \case
Hof (TypedHof _ (For _ _ (UnaryLamExpr b body))) ->
Expand All @@ -419,16 +418,13 @@ bundleFold emptyVal pair els = case els of
h:t -> (pair h tb, td + 1)
where (tb, td) = bundleFold emptyVal pair t

mkBundleTy :: [Type r n] -> (Type r n, BundleDesc)
mkBundleTy :: [Type n] -> (Type n, BundleDesc)
mkBundleTy = bundleFold UnitTy (\x y -> TyCon (ProdType [x, y]))

mkBundle :: [Atom r n] -> (Atom r n, BundleDesc)
mkBundle :: [Atom n] -> (Atom n, BundleDesc)
mkBundle = bundleFold UnitVal (\x y -> Con (ProdCon [x, y]))

freeAtomVarsList :: forall r e n. (IRRep r, HoistableE e) => e n -> [Name (AtomNameC r) n]
freeAtomVarsList = freeVarsList

freshNameM :: (Color c, EnvReader m) => NameHint -> m n (Abs (NameBinder c) (Name c) n)
freshNameM :: (EnvReader m) => NameHint -> m n (Abs NameBinder Name n)
freshNameM hint = do
scope <- toScope <$> unsafeGetEnv
Distinct <- getDistinct
Expand Down
5 changes: 2 additions & 3 deletions src/lib/DPS.hs
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ import Builder
import Core
import Imp
import CheapReduction
import IRVariants
import Name
import Subst
import PPrint
Expand All @@ -39,7 +38,7 @@ dpsPass (TopLam False piTy (LamExpr bs body)) = do
return UnitVal
dpsPass (TopLam True _ _) = error "already in destination style"

computeDPSPiTy :: PiType SimpIR i -> DestM i o (PiType SimpIR o)
computeDPSPiTy :: PiType i -> DestM i o (PiType o)
computeDPSPiTy (PiType bs resultTy) = case bs of
Empty -> do
destTy <- computeDestTy =<< dpsSubstType resultTy
Expand All @@ -57,7 +56,7 @@ type MaybeDest d n = SMaybe d (Dest n)
type MaybeResult d n = SMaybe (Not d) (SAtom n)

data DPSTag
type DestM = AtomSubstBuilder DPSTag SimpIR
type DestM = AtomSubstBuilder DPSTag

computeRepTy :: EnvReader m => SType n -> m n (SType n)
computeRepTy ty = case ty of
Expand Down
29 changes: 14 additions & 15 deletions src/lib/Generalize.hs
Original file line number Diff line number Diff line change
Expand Up @@ -14,28 +14,27 @@ import Err
import PPrint
import Types.Core
import Inference
import IRVariants
import QueryType
import Name
import Subst
import Types.Primitives
import Types.Top

generalizeIxDict :: EnvReader m => CDict n -> m n (Generalized CoreIR CDict n)
generalizeIxDict :: EnvReader m => CDict n -> m n (Generalized CDict n)
generalizeIxDict dict = liftGeneralizerM do
dict' <- sinkM dict
dictTy <- return $ getType dict'
dictTyGeneralized <- generalizeType dictTy
liftEnvReaderM $ generalizeDict dictTyGeneralized dict'
{-# INLINE generalizeIxDict #-}

generalizeArgs ::EnvReader m => CorePiType n -> [Atom CoreIR n] -> m n (Generalized CoreIR (ListE CAtom) n)
generalizeArgs ::EnvReader m => CorePiType n -> [Atom n] -> m n (Generalized (ListE CAtom) n)
generalizeArgs fTy argsTop = liftGeneralizerM $ runSubstReaderT idSubst do
PairE (CorePiType _ expls bs _) (ListE argsTop') <- sinkM $ PairE fTy (ListE argsTop)
ListE <$> go (zipAttrs expls bs) argsTop'
where
go :: Nest (WithAttrB Explicitness CBinder) i i' -> [Atom CoreIR n]
-> SubstReaderT AtomSubstVal GeneralizerM i n [Atom CoreIR n]
go :: Nest (WithAttrB Explicitness CBinder) i i' -> [Atom n]
-> SubstReaderT AtomSubstVal GeneralizerM i n [Atom n]
go (Nest (WithAttrB expl b) bs) (arg:args) = do
ty' <- substM $ binderType b
arg' <- liftSubstReaderT case (ty', expl) of
Expand All @@ -57,7 +56,7 @@ generalizeArgs fTy argsTop = liftGeneralizerM $ runSubstReaderT idSubst do

-- === generalizer monad plumbing ===

data GeneralizationEmission n l = GeneralizationEmission (Binder CoreIR n l) (Atom CoreIR n)
data GeneralizationEmission n l = GeneralizationEmission (Binder n l) (Atom n)
type GeneralizationEmissions = RNest GeneralizationEmission

newtype GeneralizerM n a = GeneralizerM {
Expand All @@ -67,7 +66,7 @@ newtype GeneralizerM n a = GeneralizerM {
liftGeneralizerM
:: EnvReader m
=> (forall l. DExt n l => GeneralizerM l (e l))
-> m n (Generalized CoreIR e n)
-> m n (Generalized e n)
liftGeneralizerM cont = do
env <- unsafeGetEnv
Distinct <- getDistinct
Expand All @@ -77,7 +76,7 @@ liftGeneralizerM cont = do
return (Abs bs e, vals)
where
-- OPTIMIZE: something not O(N^2)
hoistGeneralizationVals :: Nest GeneralizationEmission n l -> (Nest (Binder CoreIR) n l, [Atom CoreIR n])
hoistGeneralizationVals :: Nest GeneralizationEmission n l -> (Nest (Binder) n l, [Atom n])
hoistGeneralizationVals Empty = (Empty, [])
hoistGeneralizationVals (Nest (GeneralizationEmission b val) bs) = do
let (bs', vals) = hoistGeneralizationVals bs
Expand All @@ -89,7 +88,7 @@ liftGeneralizerM cont = do
{-# INLINE liftGeneralizerM #-}

-- XXX: the supplied type may be more general than the type of the atom!
emitGeneralizationParameter :: Type CoreIR n -> Atom CoreIR n -> GeneralizerM n (AtomVar CoreIR n)
emitGeneralizationParameter :: Type n -> Atom n -> GeneralizerM n (AtomVar n)
emitGeneralizationParameter ty val = GeneralizerM do
Abs b v <- return $ newName noHint
let emission = Abs (RNest REmpty (GeneralizationEmission (b:>ty) val)) v
Expand All @@ -103,7 +102,7 @@ emitGeneralizationParameter ty val = GeneralizerM do
-- === actual generalization traversal ===

-- Given a type (an Atom of type `Type`), abstracts over all data components
generalizeType :: Type CoreIR n -> GeneralizerM n (Type CoreIR n)
generalizeType :: Type n -> GeneralizerM n (Type n)
generalizeType ty = traverseTyParams ty \paramRole paramReqTy param -> case paramRole of
TypeParam -> toAtom <$> generalizeType (fromJust $ toMaybeType param)
DictParam -> toAtom <$> generalizeDict paramReqTy (fromJust $ toMaybeDict param)
Expand Down Expand Up @@ -155,13 +154,13 @@ traverseTyParams (TyCon ty) f = liftM TyCon $ getDistinct >>= \Distinct -> case

traverseRoleBinders
:: forall m n n'. EnvReader m
=> (forall l . DExt n l => ParamRole -> Type CoreIR l -> Atom CoreIR l -> m l (Atom CoreIR l))
-> Nest CBinder n n' -> [Atom CoreIR n] -> m n [Atom CoreIR n]
=> (forall l . DExt n l => ParamRole -> Type l -> Atom l -> m l (Atom l))
-> Nest CBinder n n' -> [Atom n] -> m n [Atom n]
traverseRoleBinders f allBinders allParams =
runSubstReaderT idSubst $ go allBinders allParams
where
go :: forall i i'. Nest CBinder i i' -> [Atom CoreIR n]
-> SubstReaderT AtomSubstVal m i n [Atom CoreIR n]
go :: forall i i'. Nest CBinder i i' -> [Atom n]
-> SubstReaderT AtomSubstVal m i n [Atom n]
go Empty [] = return []
go (Nest b bs) (param:params) = do
ty' <- substM $ binderType b
Expand All @@ -176,7 +175,7 @@ traverseRoleBinders f allBinders allParams =
-- === instances ===

instance GenericB GeneralizationEmission where
type RepB GeneralizationEmission = BinderP (AtomNameC CoreIR) (PairE CType CAtom)
type RepB GeneralizationEmission = BinderP (PairE CType CAtom)
fromB (GeneralizationEmission (b:>ty) x) = b :> PairE ty x
{-# INLINE fromB #-}
toB (b :> PairE ty x) = GeneralizationEmission (b:>ty) x
Expand Down
95 changes: 0 additions & 95 deletions src/lib/IRVariants.hs

This file was deleted.

Loading

0 comments on commit 0e6e527

Please sign in to comment.