Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor TopEnvFrag #861

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 21 additions & 17 deletions src/lib/Builder.hs
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,9 @@ module Builder (
emitBlock, emitDecls, BuilderEmissions, emitAtomToName,
TopBuilder (..), TopBuilderT (..), liftTopBuilderTWith, runTopBuilderT, TopBuilder2,
emitSourceMap, emitSynthCandidates, emitTopLet, emitImpFunBinding,
lookupLoadedModule, bindModule, getAllRequiredObjectFiles, extendCache,
lookupLoadedModule, registerLoadedModule, importModule, getAllRequiredObjectFiles, extendCache,
extendImpCache, queryImpCache, extendObjCache, queryObjCache, getCache, emitObjFile,
TopEnvFrag (..), emitPartialTopEnvFrag, emitLocalModuleEnv,
TopEnvFrag (..),
inlineLastDecl, fabricateEmitsEvidence, fabricateEmitsEvidenceM,
singletonBinderNest, varsAsBinderNest, typesAsBinderNest,
liftBuilder, liftEmitBuilder, makeBlock,
Expand Down Expand Up @@ -74,6 +74,7 @@ import {-# SOURCE #-} Interpreter
import LabeledItems
import Util (enumerate, restructure, transitiveClosureM, bindM2)
import Err
import Types.Core
import Core

-- === Ordinary (local) builder class ===
Expand Down Expand Up @@ -139,17 +140,14 @@ class (EnvReader m, MonadFail1 m)
=> (forall l. (Mut l, DExt n l) => m l (e l))
-> m n (Abs TopEnvFrag e n)

emitPartialTopEnvFrag :: TopBuilder m => PartialTopEnvFrag n -> m n ()
emitPartialTopEnvFrag frag = emitNamelessEnv $ TopEnvFrag emptyOutFrag frag

emitLocalModuleEnv :: TopBuilder m => ModuleEnv n -> m n ()
emitLocalModuleEnv env = emitPartialTopEnvFrag $ mempty { fragLocalModuleEnv = env }
emitIntoLocalModule :: TopBuilder m => ModuleEmissions n -> m n ()
emitIntoLocalModule me = emitNamelessEnv $ TopEnvFrag emptyOutFrag me

emitSourceMap :: TopBuilder m => SourceMap n -> m n ()
emitSourceMap sm = emitLocalModuleEnv $ mempty {envSourceMap = sm}
emitSourceMap sm = emitIntoLocalModule $ mempty { newSourceBindings = sm}

emitSynthCandidates :: TopBuilder m => SynthCandidates n -> m n ()
emitSynthCandidates sc = emitLocalModuleEnv $ mempty {envSynthCandidates = sc}
emitSynthCandidates sc = emitIntoLocalModule $ mempty { newSynthCandidates = sc}

emitTopLet :: (Mut n, TopBuilder m) => NameHint -> LetAnn -> Expr n -> m n (AtomName n)
emitTopLet hint letAnn expr = do
Expand All @@ -159,10 +157,17 @@ emitTopLet hint letAnn expr = do
emitImpFunBinding :: (Mut n, TopBuilder m) => NameHint -> ImpFunction n -> m n (ImpFunName n)
emitImpFunBinding hint f = emitBinding hint $ ImpFunBinding f

bindModule :: (Mut n, TopBuilder m) => ModuleSourceName -> ModuleName n -> m n ()
bindModule sourceName internalName = do
importModule :: (Mut n, TopBuilder m, Fallible1 m) => ModuleSourceName -> m n ()
importModule sname = do
lookupLoadedModule sname >>= \case
Nothing -> throw ModuleImportErr $ "Couldn't import " ++ pprint sname
Just name -> lookupModule name >>= \case
Module _ _ transImports _ _ _ -> emitIntoLocalModule $ ModuleEmissions (ImportStatus (S.singleton name) (S.insert name transImports)) mempty mempty mempty

registerLoadedModule :: (Mut n, TopBuilder m) => ModuleSourceName -> ModuleName n -> m n ()
registerLoadedModule sourceName internalName = do
let loaded = LoadedModules $ M.singleton sourceName internalName
emitPartialTopEnvFrag $ mempty {fragLoadedModules = loaded}
emitNamelessEnv $ TopEnvFrag (TopEnvEmissions emptyOutFrag mempty loaded) mempty

getAllRequiredObjectFiles :: EnvReader m => m n [ObjectFileName n]
getAllRequiredObjectFiles = do
Expand All @@ -176,7 +181,7 @@ lookupLoadedModule name = do
return $ M.lookup name $ fromLoadedModules loadedModules

extendCache :: TopBuilder m => Cache n -> m n ()
extendCache cache = emitPartialTopEnvFrag $ mempty {fragCache = cache}
extendCache c = emitNamelessEnv $ TopEnvFrag (TopEnvEmissions emptyOutFrag c mempty) mempty

extendImpCache :: TopBuilder m => AtomName n -> ImpFunName n -> m n ()
extendImpCache fSimp fImp =
Expand All @@ -194,7 +199,7 @@ extendObjCache fImp cfun =
emitObjFile :: (Mut n, TopBuilder m) => NameHint -> ObjectFile n -> m n (ObjectFileName n)
emitObjFile hint objFile = do
v <- emitBinding hint $ ObjectFileBinding objFile
emitLocalModuleEnv $ mempty {envObjectFiles = ObjectFiles $ S.singleton v}
emitIntoLocalModule $ mempty {newObjectFiles = ObjectFiles $ S.singleton v}
return v

queryObjCache :: EnvReader m => ImpFunName n -> m n (Maybe (CFun n))
Expand Down Expand Up @@ -228,10 +233,9 @@ instance Fallible m => TopBuilder (TopBuilderT m) where
Abs b v <- freshNameM hint
let ab = Abs (b:>binding) v
ab' <- liftEnvReaderM $ refreshAbs ab \b' v' -> do
let envFrag = toEnvFrag b'
let envFrag@(EnvFrag bindings Nothing) = toEnvFrag b'
let scs = bindingsFragToSynthCandidates envFrag
let topFrag = TopEnvFrag envFrag $
mempty { fragLocalModuleEnv = mempty { envSynthCandidates = scs} }
let topFrag = TopEnvFrag (TopEnvEmissions bindings mempty mempty) $ mempty { newSynthCandidates = scs }
return $ Abs topFrag v'
TopBuilderT $ extendInplaceT ab'

Expand Down
3 changes: 2 additions & 1 deletion src/lib/Syntax.hs
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,8 @@ module Syntax (
EnvReader (..), EnvReaderI (..), EnvReaderIT (..), runEnvReaderIT,
extendI, liftEnvReaderI, refreshAbsI, refreshLamI,
EnvExtender (..), Binding (..),
TopEnvFrag (..), ToBinding (..), PartialTopEnvFrag (..), withFreshBinders,
TopEnvFrag (..), ToBinding (..), withFreshBinders,
ModuleEmissions (..),
refreshBinders, substBinders, withFreshBinder,
withFreshLamBinder, withFreshPureLamBinder,
withFreshPiBinder, catEnvFrags,
Expand Down
18 changes: 4 additions & 14 deletions src/lib/TopLevel.hs
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,7 @@ ensureModuleLoaded moduleSourceName = do
depsRequired <- findDepsTransitively moduleSourceName
forM_ depsRequired \md -> do
evaluated <- evalPartiallyParsedUModuleCached md
bindModule (umppName md) evaluated
registerLoadedModule (umppName md) evaluated
{-# SCC ensureModuleLoaded #-}

evalSourceBlock :: (Topper m, Mut n)
Expand Down Expand Up @@ -333,22 +333,12 @@ evalPartiallyParsedUModule partiallyParsed = do
evalUModule :: (Topper m ,Mut n) => UModule -> m n (Module n)
evalUModule (UModule name _ blocks) = do
Abs topFrag UnitE <- localTopBuilder $ mapM_ (evalSourceBlock' name) blocks >> return UnitE
TopEnvFrag envFrag (PartialTopEnvFrag cache loadedModules moduleEnv) <- return topFrag
ModuleEnv (ImportStatus directDeps transDeps) sm scs objs _ <- return moduleEnv
let fragToReEmit = TopEnvFrag envFrag $ PartialTopEnvFrag cache loadedModules mempty
TopEnvFrag topEmissions moduleEmissions <- return topFrag
let ModuleEmissions (ImportStatus directDeps transDeps) sm scs objs = moduleEmissions
let fragToReEmit = TopEnvFrag topEmissions mempty
let evaluatedModule = Module name directDeps transDeps sm scs objs
emitEnv $ Abs fragToReEmit evaluatedModule

importModule :: (Mut n, TopBuilder m, Fallible1 m) => ModuleSourceName -> m n ()
importModule name = do
lookupLoadedModule name >>= \case
Nothing -> throw ModuleImportErr $ "Couldn't import " ++ pprint name
Just name' -> do
Module _ _ transImports' _ _ _ <- lookupModule name'
let importStatus = ImportStatus (S.singleton name') (S.singleton name' <> transImports')
emitLocalModuleEnv $ mempty { envImportStatus = importStatus }
{-# SCC importModule #-}

evalFile :: (Topper m, Mut n) => FilePath -> m n [(SourceBlock, Result)]
evalFile fname = evalSourceText =<< (liftIO $ readFile fname)

Expand Down
187 changes: 124 additions & 63 deletions src/lib/Types/Core.hs
Original file line number Diff line number Diff line change
Expand Up @@ -277,14 +277,19 @@ data ImportStatus (n::S) = ImportStatus
, transImports :: S.Set (ModuleName n) }
deriving (Show, Generic)

data TopEnvFrag n l = TopEnvFrag (EnvFrag n l) (PartialTopEnvFrag l)
data TopEnvFrag n l = TopEnvFrag {-# UNPACK #-} (TopEnvEmissions n l)
{-# UNPACK #-} (ModuleEmissions l)

-- This is really the type of updates to `Env`. We should probably change the
-- names to reflect that.
data PartialTopEnvFrag n = PartialTopEnvFrag
{ fragCache :: Cache n
, fragLoadedModules :: LoadedModules n
, fragLocalModuleEnv :: ModuleEnv n }
data TopEnvEmissions n l = TopEnvEmissions (RecSubstFrag Binding n l) (Cache l) (LoadedModules l)

data ModuleEmissions n where
ModuleEmissions ::
{ newImports :: ImportStatus n
-- Those never include imported values, only those emitted in the current module!
, newSourceBindings :: SourceMap n
, newSynthCandidates :: SynthCandidates n
, newObjectFiles :: ObjectFiles n
} -> ModuleEmissions n

-- TODO: we could add a lot more structure for querying by dict type, caching, etc.
-- TODO: make these `Name n` instead of `Atom n` so they're usable as cache keys.
Expand Down Expand Up @@ -1505,81 +1510,137 @@ instance ProvesExt EnvFrag
instance BindsNames EnvFrag
instance SubstB Name EnvFrag

instance GenericE PartialTopEnvFrag where
type RepE PartialTopEnvFrag = Cache
`PairE` LoadedModules
`PairE` ModuleEnv
fromE (PartialTopEnvFrag cache loaded env) = cache `PairE` loaded `PairE` env
{-# INLINE fromE #-}
toE (cache `PairE` loaded `PairE` env) = PartialTopEnvFrag cache loaded env
{-# INLINE toE #-}

instance SinkableE PartialTopEnvFrag
instance HoistableE PartialTopEnvFrag
instance AlphaEqE PartialTopEnvFrag
instance SubstE Name PartialTopEnvFrag

instance Semigroup (PartialTopEnvFrag n) where
PartialTopEnvFrag x1 x2 x3 <> PartialTopEnvFrag y1 y2 y3 =
PartialTopEnvFrag (x1<>y1) (x2<>y2) (x3<>y3)
instance GenericB TopEnvEmissions where
type RepB TopEnvEmissions = (RecSubstFrag Binding) `PairB` (LiftB (Cache `PairE` LoadedModules))
fromB (TopEnvEmissions bs c lm) = bs `PairB` (LiftB (c `PairE` lm))
toB (bs `PairB` (LiftB (c `PairE` lm))) = TopEnvEmissions bs c lm
instance SubstB Name TopEnvEmissions
instance HoistableB TopEnvEmissions
instance SinkableB TopEnvEmissions
instance ProvesExt TopEnvEmissions
instance BindsNames TopEnvEmissions

instance OutFrag TopEnvEmissions where
emptyOutFrag = TopEnvEmissions emptyOutFrag mempty mempty
{-# INLINE emptyOutFrag #-}
catOutFrags scope (TopEnvEmissions frag1 c lm) (TopEnvEmissions frag2 c' lm') =
withExtEvidence frag2 $
TopEnvEmissions (catOutFrags scope frag1 frag2) (sink c <> c') (sink lm <> lm')
{-# INLINE catOutFrags #-}

instance Monoid (PartialTopEnvFrag n) where
mempty = PartialTopEnvFrag mempty mempty mempty
mappend = (<>)
instance GenericE ModuleEmissions where
type RepE ModuleEmissions = ImportStatus
`PairE` SourceMap
`PairE` SynthCandidates
`PairE` ObjectFiles
fromE (ModuleEmissions imports sm sc obj) = imports `PairE` sm `PairE` sc `PairE` obj
toE (imports `PairE` sm `PairE` sc `PairE` obj) = ModuleEmissions imports sm sc obj
instance SinkableE ModuleEmissions
instance HoistableE ModuleEmissions
instance SubstE Name ModuleEmissions

instance Semigroup (ModuleEmissions n) where
(ModuleEmissions i sm scs obj) <> (ModuleEmissions i' sm' scs' obj') =
ModuleEmissions (i<>i') (sm<>sm') (scs<>scs') (obj<>obj')
instance Monoid (ModuleEmissions n) where
mempty = ModuleEmissions mempty mempty mempty mempty

instance GenericB TopEnvFrag where
type RepB TopEnvFrag = PairB EnvFrag (LiftB PartialTopEnvFrag)
fromB (TopEnvFrag frag1 frag2) = PairB frag1 (LiftB frag2)
toB (PairB frag1 (LiftB frag2)) = TopEnvFrag frag1 frag2

type RepB TopEnvFrag = TopEnvEmissions `PairB` (LiftB ModuleEmissions)
fromB (TopEnvFrag te me) = te `PairB` (LiftB me)
toB (te `PairB` (LiftB me)) = TopEnvFrag te me
instance SubstB Name TopEnvFrag
instance HoistableB TopEnvFrag
instance SinkableB TopEnvFrag
instance SinkableB TopEnvFrag
instance ProvesExt TopEnvFrag
instance BindsNames TopEnvFrag

instance OutFrag TopEnvFrag where
emptyOutFrag = TopEnvFrag emptyOutFrag mempty
{-# INLINE emptyOutFrag #-}
catOutFrags scope (TopEnvFrag frag1 partial1)
(TopEnvFrag frag2 partial2) =
withExtEvidence frag2 $
TopEnvFrag
(catOutFrags scope frag1 frag2)
(sink partial1 <> partial2)
catOutFrags scope (TopEnvFrag te me) (TopEnvFrag te' me') =
withExtEvidence te' $ TopEnvFrag (catOutFrags scope te te') (sink me <> me')
{-# INLINE catOutFrags #-}


--data CacheFrag (n::S) (l::S) where
--CacheFrag :: Cache n -> CacheFrag n n

--instance GenericB CacheFrag where
--type RepB CacheFrag = LiftB Cache
--fromB (CacheFrag c) = LiftB c
--toB (LiftB c) = CacheFrag c
--instance SinkableB CacheFrag
--instance ProvesExt CacheFrag
--instance BindsNames CacheFrag

--instance OutFrag CacheFrag where
--emptyOutFrag = CacheFrag mempty
--catOutFrags _ (CacheFrag a) (CacheFrag b) = CacheFrag $ a <> b

--instance ExtOutFrag any TopEnvFrag CacheFrag where
--extendOutFrag _ (TopEnvFrag f cache loaded mfrag) (CacheFrag cache') =
--TopEnvFrag f (cache <> cache') loaded mfrag

--instance ExtOutMap Env CacheFrag where
--extendOutMap env (CacheFrag cache') = Env (TopEnv defs (cache <> cache') loaded) moduleEnv
--where Env (TopEnv defs cache loaded) moduleEnv = env

--data ImportFrag (n::S) (l::S) where
--ImportFrag :: [ModuleName n] -> ImportFrag n n

--instance GenericB ImportFrag where
--type RepB ImportFrag = LiftB (ListE ModuleName)
--fromB (ImportFrag l) = LiftB (ListE l)
--toB (LiftB (ListE l)) = ImportFrag l
--instance SinkableB ImportFrag
--instance ProvesExt ImportFrag
--instance BindsNames ImportFrag

--instance OutFrag ImportFrag where
--emptyOutFrag = ImportFrag mempty
--catOutFrags _ (ImportFrag a) (ImportFrag b) = ImportFrag $ a <> b

--instance ExtOutFrag Env TopEnvFrag ImportFrag where
--extendOutFrag (TopEnvFrag bs cache loaded me) (ImportFrag importedNames) =
--TopEnvFrag bs cache loaded $ me { newImports = foldl' S.insert (newImports me) importedNames }

--instance ExtOutMap Env ImportFrag where
--extendOutMap (Env topEnv mEnv) (ImportFrag importedNames) =
--Env topEnv $ foldl' (importIntoModule topEnv) mEnv importedNames

--instance ExtOutFrag any TopEnvFrag TopEnvFrag where
--extendOutFrag _ = catOutFrags (error "shouldn't need scope")

-- XXX: unlike `ExtOutMap Env EnvFrag` instance, this once doesn't
-- extend the synthesis candidates based on the annotated let-bound names. It
-- only extends synth candidates when they're supplied explicitly.
instance ExtOutMap Env TopEnvFrag where
extendOutMap env (TopEnvFrag (EnvFrag frag _) (PartialTopEnvFrag cache' loaded' mEnv')) = result
extendOutMap env (TopEnvFrag (TopEnvEmissions frag cache' loaded') mEnv') = Env newTopEnv newModuleEnv
where
Env (TopEnv defs cache loaded) mEnv = env
result = Env newTopEnv newModuleEnv

newTopEnv = withExtEvidence frag $ TopEnv
(defs `extendRecSubst` frag)
(sink cache <> cache')
(sink loaded <> loaded')

newModuleEnv =
ModuleEnv
(imports <> imports')
(sm <> sm' <> newImportedSM)
(scs <> scs' <> newImportedSC)
(objs <> objs' <> newImportedObj)
(effs <> effs')
where
ModuleEnv imports sm scs objs effs = withExtEvidence frag $ sink mEnv
ModuleEnv imports' sm' scs' objs' effs' = mEnv'
newDirectImports = S.difference (directImports imports') (directImports imports)
newTransImports = S.difference (transImports imports') (transImports imports)
newImportedSM = flip foldMap newDirectImports $ moduleExports . lookupModulePure
newImportedSC = flip foldMap newTransImports $ moduleSynthCandidates . lookupModulePure
newImportedObj = flip foldMap newTransImports $ moduleObjectFiles . lookupModulePure

lookupModulePure v = case lookupEnvPure (Env newTopEnv mempty) v of ModuleBinding m -> m
newTopEnv = withExtEvidence frag $ TopEnv (defs `extendRecSubst` frag) (sink cache <> cache') (sink loaded <> loaded')

ModuleEmissions imports' sm' scs' objs' = mEnv'
ModuleEnv imports sm scs objs effs = withExtEvidence frag $
foldl' (importIntoModule newTopEnv) (sink mEnv) $ directImports imports'
newModuleEnv = ModuleEnv imports (sm <> sm') (scs <> scs') (objs <> objs') effs

importIntoModule :: TopEnv n -> ModuleEnv n -> ModuleName n -> ModuleEnv n
importIntoModule topEnv menv@(ModuleEnv imports sm scs objs effs) importName =
case importName `S.member` (directImports imports) of
True -> menv
False -> ModuleEnv (imports <> imports') (sm <> newImportedSM) (scs <> newImportedSCS) (objs <> newImportedObj) effs
where
Module _ _ importTransImports _ _ _ = lookupModulePure importName
imports' = ImportStatus (S.singleton importName) (S.singleton importName <> importTransImports)
newDirectImports = S.difference (directImports imports') (directImports imports)
newTransImports = S.difference (transImports imports') (transImports imports)
newImportedSM = flip foldMap newDirectImports $ moduleExports . lookupModulePure
newImportedSCS = flip foldMap newTransImports $ moduleSynthCandidates . lookupModulePure
newImportedObj = flip foldMap newTransImports $ moduleObjectFiles . lookupModulePure

lookupModulePure v = case lookupEnvPure (Env topEnv mempty) v of ModuleBinding m -> m

lookupEnvPure :: Color c => Env n -> Name c n -> Binding c n
lookupEnvPure env v = lookupTerminalSubstFrag (fromRecSubst $ envDefs $ topEnv env) v
Expand Down