content: -- Copyright 2022 Google LLC
  --
  -- Use of this source code is governed by a BSD-style
  -- license that can be found in the LICENSE file or at
  -- https://developers.google.com/open-source/licenses/bsd
  
  {-# LANGUAGE UndecidableInstances #-}
  {-# LANGUAGE StrictData #-}
  
  -- Top-level data types
  
  module Types.Top where
  
  import Data.Functor ((<&>))
  import Data.Hashable
  import Data.Text.Prettyprint.Doc
  import qualified Data.Map.Strict       as M
  import qualified Data.Set              as S
  
  import GHC.Generics (Generic (..))
  import Data.Store (Store (..))
  import Foreign.Ptr
  
  import Name
  import Util (FileHash, SnocList (..))
  import IRVariants
  import PPrint
  
  import Types.Primitives
  import Types.Core
  import Types.Source
  import Types.Imp
  
  type TopBlock = TopLam -- used for nullary lambda
  type IsDestLam = Bool
  data TopLam (r::IR) (n::S) = TopLam IsDestLam (PiType r n) (LamExpr r n)
       deriving (Show, Generic)
  type STopLam = TopLam SimpIR
  type CTopLam = TopLam CoreIR
  
  data EvalStatus a = Waiting | Running | Finished a
       deriving (Show, Eq, Ord, Generic, Functor, Foldable, Traversable)
  type TopFunEvalStatus n = EvalStatus (TopFunLowerings n)
  
  data TopFun (n::S) =
     DexTopFun (TopFunDef n) (TopLam SimpIR n) (TopFunEvalStatus n)
   | FFITopFun String IFunType
     deriving (Show, Generic)
  
  data TopFunDef (n::S) =
     Specialization       (SpecializationSpec n)
   | LinearizationPrimal  (LinearizationSpec n)
     -- Tangent functions all take some number of nonlinear args, then a *single*
     -- linear arg. This is so that transposition can be an involution - you apply
     -- it twice and you get back to the original function.
   | LinearizationTangent (LinearizationSpec n)
     deriving (Show, Generic)
  
  newtype TopFunLowerings (n::S) = TopFunLowerings
    { topFunObjCode :: FunObjCodeName n } -- TODO: add optimized, imp etc. as needed
    deriving (Show, Generic, SinkableE, HoistableE, RenameE, AlphaEqE, AlphaHashableE, Pretty)
  
  data AtomBinding (r::IR) (n::S) where
   LetBound     :: DeclBinding r  n  -> AtomBinding r n
   MiscBound    :: Type        r  n  -> AtomBinding r n
   TopDataBound :: RepVal n          -> AtomBinding SimpIR n
   SolverBound  :: SolverBinding n   -> AtomBinding CoreIR n
   NoinlineFun  :: CType n -> CAtom n -> AtomBinding CoreIR n
   FFIFunBound  :: CorePiType n -> TopFunName n -> AtomBinding CoreIR n
  
  deriving instance IRRep r => Show (AtomBinding r n)
  deriving via WrapE (AtomBinding r) n instance IRRep r => Generic (AtomBinding r n)
  
  data SolverBinding (n::S) =
     InfVarBound (CType n)
   | SkolemBound (CType n)
   | DictBound   (CType n)
     deriving (Show, Generic)
  
  -- TODO: Use an IntMap
  newtype CustomRules (n::S) =
    CustomRules { customRulesMap :: M.Map (AtomName CoreIR n) (AtomRules n) }
    deriving (Semigroup, Monoid, Store)
  data AtomRules (n::S) =
    -- number of implicit args, number of explicit args, linearization function
    CustomLinearize Int Int SymbolicZeros (CAtom n)
    deriving (Generic)
  
  -- === envs and modules ===
  
  -- `ModuleEnv` contains data that only makes sense in the context of evaluating
  -- a particular module. `TopEnv` contains everything that makes sense "between"
  -- evaluating modules.
  data Env n = Env
    { topEnv    :: {-# UNPACK #-} TopEnv n
    , moduleEnv :: {-# UNPACK #-} ModuleEnv n }
    deriving (Generic)
  
  newtype EnvFrag (n::S) (l::S) = EnvFrag (RecSubstFrag Binding n l)
          deriving (OutFrag)
  
  data TopEnv (n::S) = TopEnv
    { envDefs          :: RecSubst Binding n
    , envCustomRules   :: CustomRules n
    , envCache         :: Cache n
    , envLoadedModules :: LoadedModules n
    , envLoadedObjects :: LoadedObjects n }
    deriving (Generic)
  
  data SerializedEnv n = SerializedEnv
    { serializedEnvDefs        :: RecSubst Binding n
    , serializedEnvCustomRules :: CustomRules n
    , serializedEnvCache       :: Cache n }
    deriving (Generic)
  
  -- TODO: consider splitting this further into `ModuleEnv` (the env that's
  -- relevant between top-level decls) and `LocalEnv` (the additional parts of the
  -- env that's relevant under a lambda binder). Unlike the Top/Module
  -- distinction, there's some overlap. For example, instances can be defined at
  -- both the module-level and local level. Similarly, if we start allowing
  -- top-level effects in `Main` then we'll have module-level effects and local
  -- effects.
  data ModuleEnv (n::S) = ModuleEnv
    { envImportStatus    :: ImportStatus n
    , envSourceMap       :: SourceMap n
    , envSynthCandidates :: SynthCandidates n }
    deriving (Generic)
  
  data Module (n::S) = Module
    { moduleSourceName :: ModuleSourceName
    , moduleDirectDeps :: S.Set (ModuleName n)
    , moduleTransDeps  :: S.Set (ModuleName n)  -- XXX: doesn't include the module itself
    , moduleExports    :: SourceMap n
      -- these are just the synth candidates required by this
      -- module by itself. We'll usually also need those required by the module's
      -- (transitive) dependencies, which must be looked up separately.
    , moduleSynthCandidates :: SynthCandidates n }
    deriving (Show, Generic)
  
  data LoadedModules (n::S) = LoadedModules
    { fromLoadedModules   :: M.Map ModuleSourceName (ModuleName n)}
    deriving (Show, Generic)
  
  emptyModuleEnv :: ModuleEnv n
  emptyModuleEnv = ModuleEnv emptyImportStatus (SourceMap mempty) mempty
  
  emptyLoadedModules :: LoadedModules n
  emptyLoadedModules = LoadedModules mempty
  
  data LoadedObjects (n::S) = LoadedObjects
    -- the pointer points to the actual runtime function
    { fromLoadedObjects :: M.Map (FunObjCodeName n) NativeFunction}
    deriving (Show, Generic)
  
  emptyLoadedObjects :: LoadedObjects n
  emptyLoadedObjects = LoadedObjects mempty
  
  data ImportStatus (n::S) = ImportStatus
    { directImports :: S.Set (ModuleName n)
      -- XXX: This are cached for efficiency. It's derivable from `directImports`.
    , transImports           :: S.Set (ModuleName n) }
    deriving (Show, Generic)
  
  data TopEnvFrag n l = TopEnvFrag (EnvFrag n l) (ModuleEnv l) (SnocList (TopEnvUpdate l))
  
  data TopEnvUpdate n =
     ExtendCache              (Cache n)
   | AddCustomRule            (CAtomName n) (AtomRules n)
   | UpdateLoadedModules      ModuleSourceName (ModuleName n)
   | UpdateLoadedObjects      (FunObjCodeName n) NativeFunction
   | FinishDictSpecialization (SpecDictName n) [TopLam SimpIR n]
   | LowerDictSpecialization  (SpecDictName n) [TopLam SimpIR n]
   | UpdateTopFunEvalStatus   (TopFunName n) (TopFunEvalStatus n)
   | UpdateInstanceDef        (InstanceName n) (InstanceDef n)
   | UpdateTyConDef           (TyConName n) (TyConDef n)
   | UpdateFieldDef           (TyConName n) SourceName (CAtomName n)
  
  -- TODO: we could add a lot more structure for querying by dict type, caching, etc.
  data SynthCandidates n = SynthCandidates
    { instanceDicts     :: M.Map (ClassName n) [InstanceName n]
    , ixInstances       :: [InstanceName n] }
    deriving (Show, Generic)
  
  emptyImportStatus :: ImportStatus n
  emptyImportStatus = ImportStatus mempty mempty
  
  -- TODO: figure out the additional top-level context we need -- backend, other
  -- compiler flags etc. We can have a map from those to this.
  
  data Cache (n::S) = Cache
    { specializationCache :: EMap SpecializationSpec TopFunName n
    , ixDictCache :: EMap AbsDict SpecDictName n
    , linearizationCache :: EMap LinearizationSpec (PairE TopFunName TopFunName) n
    , transpositionCache :: EMap TopFunName TopFunName n
      -- This is memoizing `parseAndGetDeps :: Text -> [ModuleSourceName]`. But we
      -- only want to store one entry per module name as a simple cache eviction
      -- policy, so we store it keyed on the module name, with the text hash for
      -- the validity check.
    , parsedDeps :: M.Map ModuleSourceName (FileHash, [ModuleSourceName])
    , moduleEvaluations :: M.Map ModuleSourceName ((FileHash, [ModuleName n]), ModuleName n)
    } deriving (Show, Generic)
  
  -- === runtime function and variable representations ===
  
  type RuntimeEnv = DynamicVarKeyPtrs
  
  type DexDestructor = FunPtr (IO ())
  
  data NativeFunction = NativeFunction
    { nativeFunPtr      :: FunPtr ()
    , nativeFunTeardown :: IO () }
  
  instance Show NativeFunction where
    show _ = "<native function>"
  
  -- Holds pointers to thread-local storage used to simulate dynamically scoped
  -- variables, such as the output stream file descriptor.
  type DynamicVarKeyPtrs = [(DynamicVar, Ptr ())]
  
  data DynamicVar = OutStreamDyvar -- TODO: add others as needed
                    deriving (Enum, Bounded)
  
  dynamicVarCName :: DynamicVar -> String
  dynamicVarCName OutStreamDyvar = "dex_out_stream_dyvar"
  
  dynamicVarLinkMap :: DynamicVarKeyPtrs -> [(String, Ptr ())]
  dynamicVarLinkMap dyvars = dyvars <&> \(v, ptr) -> (dynamicVarCName v, ptr)
  
  -- === Specialization and generalization ===
  
  type Generalized (r::IR) (e::E) (n::S) = (Abstracted r e n, [Atom r n])
  type Abstracted (r::IR) (e::E) = Abs (Nest (Binder r)) e
  type AbsDict = Abstracted CoreIR (Dict CoreIR)
  
  data SpecializedDictDef n =
     SpecializedDict
       (AbsDict n)
       -- Methods (thunked if nullary), if they're available.
       -- We create specialized dict names during simplification, but we don't
       -- actually simplify/lower them until we return to TopLevel
       (Maybe [TopLam SimpIR n])
     deriving (Show, Generic)
  
  -- TODO: extend with AD-oriented specializations, backend-specific specializations etc.
  data SpecializationSpec (n::S) =
     AppSpecialization (AtomVar CoreIR n) (Abstracted CoreIR (ListE CAtom) n)
     deriving (Show, Generic)
  
  type Active = Bool
  data LinearizationSpec (n::S) = LinearizationSpec (TopFunName n) [Active]
       deriving (Show, Generic)
  
  -- === bindings - static information we carry about a lexical scope ===
  
  -- TODO: consider making this an open union via a typeable-like class
  data Binding (c::C) (n::S) where
    AtomNameBinding   :: AtomBinding r n                -> Binding (AtomNameC r)   n
    TyConBinding      :: Maybe (TyConDef n) -> DotMethods n -> Binding TyConNameC      n
    DataConBinding    :: TyConName n -> Int             -> Binding DataConNameC    n
    ClassBinding      :: ClassDef n                     -> Binding ClassNameC      n
    InstanceBinding   :: InstanceDef n -> CorePiType n  -> Binding InstanceNameC   n
    MethodBinding     :: ClassName n   -> Int           -> Binding MethodNameC     n
    TopFunBinding     :: TopFun n                       -> Binding TopFunNameC     n
    FunObjCodeBinding :: CFunction n                    -> Binding FunObjCodeNameC n
    ModuleBinding     :: Module n                       -> Binding ModuleNameC     n
    -- TODO: add a case for abstracted pointers, as used in `ClosedImpFunction`
    PtrBinding        :: PtrType -> PtrLitVal           -> Binding PtrNameC        n
    SpecializedDictBinding :: SpecializedDictDef n      -> Binding SpecializedDictNameC n
    ImpNameBinding    :: BaseType                       -> Binding ImpNameC n
  
  -- === ToBinding ===
  
  atomBindingToBinding :: AtomBinding r n -> Binding (AtomNameC r) n
  atomBindingToBinding b = AtomNameBinding b
  
  bindingToAtomBinding :: Binding (AtomNameC r) n -> AtomBinding r n
  bindingToAtomBinding (AtomNameBinding b) = b
  
  class (RenameE     e, SinkableE e) => ToBinding (e::E) (c::C) | e -> c where
    toBinding :: e n -> Binding c n
  
  instance Color c => ToBinding (Binding c) c where
    toBinding = id
  
  instance IRRep r => ToBinding (AtomBinding r) (AtomNameC r) where
    toBinding = atomBindingToBinding
  
  instance IRRep r => ToBinding (DeclBinding r) (AtomNameC r) where
    toBinding = toBinding . LetBound
  
  instance IRRep r => ToBinding (Type r) (AtomNameC r) where
    toBinding = toBinding . MiscBound
  
  instance ToBinding SolverBinding (AtomNameC CoreIR) where
    toBinding = toBinding . SolverBound
  
  instance IRRep r => ToBinding (IxType r) (AtomNameC r) where
    toBinding (IxType t _) = toBinding t
  
  instance (ToBinding e1 c, ToBinding e2 c) => ToBinding (EitherE e1 e2) c where
    toBinding (LeftE  e) = toBinding e
    toBinding (RightE e) = toBinding e
  
  instance ToBindersAbs (TopLam r) (Expr r) r where
    toAbs (TopLam _ _ lam) = toAbs lam
  
  -- === GenericE, GenericB ===
  
  instance GenericE SpecializedDictDef where
    type RepE SpecializedDictDef = AbsDict `PairE` MaybeE (ListE (TopLam SimpIR))
    fromE (SpecializedDict ab methods) = ab `PairE` methods'
      where methods' = case methods of Just xs -> LeftE (ListE xs)
                                       Nothing -> RightE UnitE
    {-# INLINE fromE #-}
    toE   (ab `PairE` methods) = SpecializedDict ab methods'
      where methods' = case methods of LeftE (ListE xs) -> Just xs
                                       RightE UnitE     -> Nothing
    {-# INLINE toE #-}
  
  instance SinkableE      SpecializedDictDef
  instance HoistableE     SpecializedDictDef
  instance AlphaEqE       SpecializedDictDef
  instance AlphaHashableE SpecializedDictDef
  instance RenameE        SpecializedDictDef
  
  instance HasScope Env where
    toScope = toScope . envDefs . topEnv
  
  instance OutMap Env where
    emptyOutMap =
      Env (TopEnv (RecSubst emptyInFrag) mempty mempty emptyLoadedModules emptyLoadedObjects)
          emptyModuleEnv
    {-# INLINE emptyOutMap #-}
  
  instance ExtOutMap Env (RecSubstFrag Binding)  where
    -- TODO: We might want to reorganize this struct to make this
    -- do less explicit sinking etc. It's a hot operation!
    extendOutMap (Env (TopEnv defs rules cache loadedM loadedO) moduleEnv) frag =
      withExtEvidence frag $ Env
        (TopEnv
          (defs  `extendRecSubst` frag)
          (sink rules)
          (sink cache)
          (sink loadedM)
          (sink loadedO))
        (sink moduleEnv)
    {-# INLINE extendOutMap #-}
  
  instance ExtOutMap Env EnvFrag where
    extendOutMap = extendEnv
    {-# INLINE extendOutMap #-}
  
  extendEnv :: Distinct l => Env n -> EnvFrag n l -> Env l
  extendEnv env (EnvFrag newEnv) = do
    case extendOutMap env newEnv of
      Env envTop (ModuleEnv imports sm scs) -> do
        Env envTop (ModuleEnv imports sm scs)
  {-# NOINLINE [1] extendEnv #-}
  
  
  instance GenericE AtomRules where
    type RepE AtomRules = (LiftE (Int, Int, SymbolicZeros)) `PairE` CAtom
    fromE (CustomLinearize ni ne sz a) = LiftE (ni, ne, sz) `PairE` a
    toE (LiftE (ni, ne, sz) `PairE` a) = CustomLinearize ni ne sz a
  instance SinkableE AtomRules
  instance HoistableE AtomRules
  instance AlphaEqE AtomRules
  instance RenameE     AtomRules
  
  instance GenericE CustomRules where
    type RepE CustomRules = ListE (PairE (AtomName CoreIR) AtomRules)
    fromE (CustomRules m) = ListE $ toPairE <$> M.toList m
    toE (ListE l) = CustomRules $ M.fromList $ fromPairE <$> l
  instance SinkableE CustomRules
  instance HoistableE CustomRules
  instance AlphaEqE CustomRules
  instance RenameE     CustomRules
  
  instance GenericE Cache where
    type RepE Cache =
              EMap SpecializationSpec TopFunName
      `PairE` EMap AbsDict SpecDictName
      `PairE` EMap LinearizationSpec (PairE TopFunName TopFunName)
      `PairE` EMap TopFunName TopFunName
      `PairE` LiftE (M.Map ModuleSourceName (FileHash, [ModuleSourceName]))
      `PairE` ListE (        LiftE ModuleSourceName
                     `PairE` LiftE FileHash
                     `PairE` ListE ModuleName
                     `PairE` ModuleName)
    fromE (Cache x y z w parseCache evalCache) =
      x `PairE` y `PairE` z `PairE` w `PairE` LiftE parseCache `PairE`
        ListE [LiftE sourceName `PairE` LiftE hashVal `PairE` ListE deps `PairE` result
               | (sourceName, ((hashVal, deps), result)) <- M.toList evalCache ]
    {-# INLINE fromE #-}
    toE   (x `PairE` y `PairE` z `PairE` w `PairE` LiftE parseCache `PairE` ListE evalCache) =
      Cache x y z w parseCache
        (M.fromList
         [(sourceName, ((hashVal, deps), result))
         | LiftE sourceName `PairE` LiftE hashVal `PairE` ListE deps `PairE` result
            <- evalCache])
    {-# INLINE toE #-}
  
  instance SinkableE  Cache
  instance HoistableE Cache
  instance AlphaEqE   Cache
  instance RenameE     Cache
  instance Store (Cache n)
  
  instance Monoid (Cache n) where
    mempty = Cache mempty mempty mempty mempty mempty mempty
    mappend = (<>)
  
  instance Semigroup (Cache n) where
    -- right-biased instead of left-biased
    Cache x1 x2 x3 x4 x5 x6 <> Cache y1 y2 y3 y4 y5 y6 =
      Cache (y1<>x1) (y2<>x2) (y3<>x3) (y4<>x4) (x5<>y5) (x6<>y6)
  
  
  instance GenericE SynthCandidates where
    type RepE SynthCandidates = ListE (PairE ClassName (ListE InstanceName))
                        `PairE` ListE InstanceName
    fromE (SynthCandidates xs ys) = ListE xs' `PairE` ListE ys
      where xs' = map (\(k,vs) -> PairE k (ListE vs)) (M.toList xs)
    {-# INLINE fromE #-}
    toE (ListE xs `PairE` ListE ys) = SynthCandidates xs' ys
      where xs' = M.fromList $ map (\(PairE k (ListE vs)) -> (k,vs)) xs
    {-# INLINE toE #-}
  
  instance SinkableE      SynthCandidates
  instance HoistableE     SynthCandidates
  instance AlphaEqE       SynthCandidates
  instance AlphaHashableE SynthCandidates
  instance RenameE        SynthCandidates
  
  instance IRRep r => GenericE (AtomBinding r) where
    type RepE (AtomBinding r) =
       EitherE2 (EitherE3
        (DeclBinding   r)              -- LetBound
        (Type          r)              -- MiscBound
        (WhenCore r SolverBinding)     -- SolverBound
       ) (EitherE3
        (WhenCore r (PairE CType CAtom))               -- NoinlineFun
        (WhenSimp r RepVal)                            -- TopDataBound
        (WhenCore r (CorePiType `PairE` TopFunName))   -- FFIFunBound
       )
  
    fromE = \case
      LetBound    x -> Case0 $ Case0 x
      MiscBound   x -> Case0 $ Case1 x
      SolverBound x -> Case0 $ Case2 $ WhenIRE x
      NoinlineFun t x -> Case1 $ Case0 $ WhenIRE $ PairE t x
      TopDataBound repVal -> Case1 $ Case1 $ WhenIRE repVal
      FFIFunBound ty v    -> Case1 $ Case2 $ WhenIRE $ ty `PairE` v
    {-# INLINE fromE #-}
  
    toE = \case
      Case0 x' -> case x' of
        Case0 x         -> LetBound x
        Case1 x           -> MiscBound x
        Case2 (WhenIRE x) -> SolverBound x
        _ -> error "impossible"
      Case1 x' -> case x' of
        Case0 (WhenIRE (PairE t x)) -> NoinlineFun t x
        Case1 (WhenIRE repVal)                         -> TopDataBound repVal
        Case2 (WhenIRE (ty `PairE` v))                 -> FFIFunBound ty v
        _ -> error "impossible"
      _ -> error "impossible"
    {-# INLINE toE #-}
  
  
  instance IRRep r => SinkableE   (AtomBinding r)
  instance IRRep r => HoistableE  (AtomBinding r)
  instance IRRep r => RenameE     (AtomBinding r)
  instance IRRep r => AlphaEqE       (AtomBinding r)
  instance IRRep r => AlphaHashableE (AtomBinding r)
  
  instance GenericE TopFunDef where
    type RepE TopFunDef = EitherE3 SpecializationSpec LinearizationSpec LinearizationSpec
    fromE = \case
      Specialization       s -> Case0 s
      LinearizationPrimal  s -> Case1 s
      LinearizationTangent s -> Case2 s
    {-# INLINE fromE #-}
    toE = \case
      Case0 s ->   Specialization       s
      Case1 s ->   LinearizationPrimal  s
      Case2 s ->   LinearizationTangent s
      _ -> error "impossible"
    {-# INLINE toE #-}
  
  instance SinkableE      TopFunDef
  instance HoistableE     TopFunDef
  instance RenameE        TopFunDef
  instance AlphaEqE       TopFunDef
  instance AlphaHashableE TopFunDef
  
  instance IRRep r => GenericE (TopLam r) where
    type RepE (TopLam r) = LiftE Bool `PairE` PiType r `PairE` LamExpr r
    fromE (TopLam d x y) = LiftE d `PairE` x `PairE` y
    {-# INLINE fromE #-}
    toE (LiftE d `PairE` x `PairE` y) = TopLam d x y
    {-# INLINE toE #-}
  
  instance IRRep r => SinkableE      (TopLam r)
  instance IRRep r => HoistableE     (TopLam r)
  instance IRRep r => RenameE        (TopLam r)
  instance IRRep r => AlphaEqE       (TopLam r)
  instance IRRep r => AlphaHashableE (TopLam r)
  
  instance GenericE TopFun where
    type RepE TopFun = EitherE
          (TopFunDef `PairE` TopLam SimpIR `PairE` ComposeE EvalStatus TopFunLowerings)
          (LiftE (String, IFunType))
    fromE = \case
      DexTopFun def lam status -> LeftE (def `PairE` lam `PairE` ComposeE status)
      FFITopFun name ty    -> RightE (LiftE (name, ty))
    {-# INLINE fromE #-}
    toE = \case
      LeftE  (def `PairE` lam `PairE` ComposeE status) -> DexTopFun def lam status
      RightE (LiftE (name, ty))            -> FFITopFun name ty
    {-# INLINE toE #-}
  
  instance SinkableE      TopFun
  instance HoistableE     TopFun
  instance RenameE        TopFun
  instance AlphaEqE       TopFun
  instance AlphaHashableE TopFun
  
  instance GenericE SpecializationSpec where
    type RepE SpecializationSpec =
           PairE (AtomVar CoreIR) (Abs (Nest (Binder CoreIR)) (ListE CAtom))
    fromE (AppSpecialization fname (Abs bs args)) = PairE fname (Abs bs args)
    {-# INLINE fromE #-}
    toE   (PairE fname (Abs bs args)) = AppSpecialization fname (Abs bs args)
    {-# INLINE toE #-}
  
  instance HasNameHint (SpecializationSpec n) where
    getNameHint (AppSpecialization f _) = getNameHint f
  
  instance SinkableE      SpecializationSpec
  instance HoistableE     SpecializationSpec
  instance RenameE        SpecializationSpec
  instance AlphaEqE       SpecializationSpec
  instance AlphaHashableE SpecializationSpec
  
  instance GenericE LinearizationSpec where
    type RepE LinearizationSpec = PairE TopFunName (LiftE [Active])
    fromE (LinearizationSpec fname actives) = PairE fname (LiftE actives)
    {-# INLINE fromE #-}
    toE   (PairE fname (LiftE actives)) = LinearizationSpec fname actives
    {-# INLINE toE #-}
  
  instance SinkableE      LinearizationSpec
  instance HoistableE     LinearizationSpec
  instance RenameE        LinearizationSpec
  instance AlphaEqE       LinearizationSpec
  instance AlphaHashableE LinearizationSpec
  
  instance GenericE SolverBinding where
    type RepE SolverBinding = EitherE3
                                    CType
                                    CType
                                    CType
    fromE = \case
      InfVarBound ty -> Case0 ty
      SkolemBound ty -> Case1 ty
      DictBound   ty -> Case2 ty
    {-# INLINE fromE #-}
  
    toE = \case
      Case0 ty -> InfVarBound ty
      Case1 ty -> SkolemBound ty
      Case2 ty -> DictBound   ty
      _ -> error "impossible"
    {-# INLINE toE #-}
  
  instance SinkableE   SolverBinding
  instance HoistableE  SolverBinding
  instance RenameE     SolverBinding
  instance AlphaEqE       SolverBinding
  instance AlphaHashableE SolverBinding
  
  instance GenericE (Binding c) where
    type RepE (Binding c) =
      EitherE3
        (EitherE6
            (WhenAtomName        c AtomBinding)
            (WhenC TyConNameC    c (MaybeE TyConDef `PairE` DotMethods))
            (WhenC DataConNameC  c (TyConName `PairE` LiftE Int))
            (WhenC ClassNameC    c (ClassDef))
            (WhenC InstanceNameC c (InstanceDef `PairE` CorePiType))
            (WhenC MethodNameC   c (ClassName `PairE` LiftE Int)))
        (EitherE4
            (WhenC TopFunNameC     c (TopFun))
            (WhenC FunObjCodeNameC c (CFunction))
            (WhenC ModuleNameC     c (Module))
            (WhenC PtrNameC        c (LiftE (PtrType, PtrLitVal))))
        (EitherE2
            (WhenC SpecializedDictNameC c (SpecializedDictDef))
            (WhenC ImpNameC             c (LiftE BaseType)))
  
    fromE = \case
      AtomNameBinding   binding           -> Case0 $ Case0 $ WhenAtomName binding
      TyConBinding      dataDef methods   -> Case0 $ Case1 $ WhenC $ toMaybeE dataDef `PairE` methods
      DataConBinding    dataDefName idx   -> Case0 $ Case2 $ WhenC $ dataDefName `PairE` LiftE idx
      ClassBinding      classDef          -> Case0 $ Case3 $ WhenC $ classDef
      InstanceBinding   instanceDef ty    -> Case0 $ Case4 $ WhenC $ instanceDef `PairE` ty
      MethodBinding     className idx     -> Case0 $ Case5 $ WhenC $ className `PairE` LiftE idx
      TopFunBinding     fun               -> Case1 $ Case0 $ WhenC $ fun
      FunObjCodeBinding cFun              -> Case1 $ Case1 $ WhenC $ cFun
      ModuleBinding m                     -> Case1 $ Case2 $ WhenC $ m
      PtrBinding ty p                     -> Case1 $ Case3 $ WhenC $ LiftE (ty,p)
      SpecializedDictBinding def          -> Case2 $ Case0 $ WhenC $ def
      ImpNameBinding ty                   -> Case2 $ Case1 $ WhenC $ LiftE ty
    {-# INLINE fromE #-}
  
    toE = \case
      Case0 (Case0 (WhenAtomName binding))           -> AtomNameBinding   binding
      Case0 (Case1 (WhenC (def `PairE` methods)))    -> TyConBinding      (fromMaybeE def) methods
      Case0 (Case2 (WhenC (n `PairE` LiftE idx)))    -> DataConBinding    n idx
      Case0 (Case3 (WhenC (classDef)))               -> ClassBinding      classDef
      Case0 (Case4 (WhenC (instanceDef `PairE` ty))) -> InstanceBinding   instanceDef ty
      Case0 (Case5 (WhenC ((n `PairE` LiftE i))))    -> MethodBinding     n i
      Case1 (Case0 (WhenC (fun)))                    -> TopFunBinding     fun
      Case1 (Case1 (WhenC (f)))                      -> FunObjCodeBinding f
      Case1 (Case2 (WhenC (m)))                      -> ModuleBinding     m
      Case1 (Case3 (WhenC ((LiftE (ty,p)))))         -> PtrBinding        ty p
      Case2 (Case0 (WhenC (def)))                    -> SpecializedDictBinding def
      Case2 (Case1 (WhenC ((LiftE ty))))             -> ImpNameBinding    ty
      _ -> error "impossible"
    {-# INLINE toE #-}
  
  deriving via WrapE (Binding c) n instance Generic (Binding c n)
  instance SinkableV         Binding
  instance HoistableV        Binding
  instance RenameV           Binding
  instance Color c => SinkableE   (Binding c)
  instance Color c => HoistableE  (Binding c)
  instance Color c => RenameE     (Binding c)
  
  instance Semigroup (SynthCandidates n) where
    SynthCandidates xs ys <> SynthCandidates xs' ys' =
      SynthCandidates (M.unionWith (<>) xs xs') (ys <> ys')
  
  instance Monoid (SynthCandidates n) where
    mempty = SynthCandidates mempty mempty
  
  instance GenericB EnvFrag where
    type RepB EnvFrag = RecSubstFrag Binding
    fromB (EnvFrag frag) = frag
    toB   frag = EnvFrag frag
  
  instance SinkableB   EnvFrag
  instance HoistableB  EnvFrag
  instance ProvesExt   EnvFrag
  instance BindsNames  EnvFrag
  instance RenameB     EnvFrag
  
  instance GenericE TopEnvUpdate where
    type RepE TopEnvUpdate = EitherE2 (
        EitherE4
      {- ExtendCache -}              Cache
      {- AddCustomRule -}            (CAtomName `PairE` AtomRules)
      {- UpdateLoadedModules -}      (LiftE ModuleSourceName `PairE` ModuleName)
      {- UpdateLoadedObjects -}      (FunObjCodeName `PairE` LiftE NativeFunction)
        ) ( EitherE6
      {- FinishDictSpecialization -} (SpecDictName `PairE` ListE (TopLam SimpIR))
      {- LowerDictSpecialization -}  (SpecDictName `PairE` ListE (TopLam SimpIR))
      {- UpdateTopFunEvalStatus -}   (TopFunName `PairE` ComposeE EvalStatus TopFunLowerings)
      {- UpdateInstanceDef -}        (InstanceName `PairE` InstanceDef)
      {- UpdateTyConDef -}           (TyConName `PairE` TyConDef)
      {- UpdateFieldDef -}           (TyConName `PairE` LiftE SourceName `PairE` CAtomName)
          )
    fromE = \case
      ExtendCache x                -> Case0 $ Case0 x
      AddCustomRule x y            -> Case0 $ Case1 (x `PairE` y)
      UpdateLoadedModules x y      -> Case0 $ Case2 (LiftE x `PairE` y)
      UpdateLoadedObjects x y      -> Case0 $ Case3 (x `PairE` LiftE y)
      FinishDictSpecialization x y -> Case1 $ Case0 (x `PairE` ListE y)
      LowerDictSpecialization x y  -> Case1 $ Case1 (x `PairE` ListE y)
      UpdateTopFunEvalStatus x y   -> Case1 $ Case2 (x `PairE` ComposeE y)
      UpdateInstanceDef x y        -> Case1 $ Case3 (x `PairE` y)
      UpdateTyConDef x y           -> Case1 $ Case4 (x `PairE` y)
      UpdateFieldDef x y z         -> Case1 $ Case5 (x `PairE` LiftE y `PairE` z)
  
    toE = \case
      Case0 e -> case e of
        Case0 x                   -> ExtendCache x
        Case1 (x `PairE` y)       -> AddCustomRule x y
        Case2 (LiftE x `PairE` y) -> UpdateLoadedModules x y
        Case3 (x `PairE` LiftE y) -> UpdateLoadedObjects x y
        _ -> error "impossible"
      Case1 e -> case e of
        Case0 (x `PairE` ListE y)           -> FinishDictSpecialization x y
        Case1 (x `PairE` ListE y)           -> LowerDictSpecialization x y
        Case2 (x `PairE` ComposeE y)        -> UpdateTopFunEvalStatus x y
        Case3 (x `PairE` y)                 -> UpdateInstanceDef x y
        Case4 (x `PairE` y)                 -> UpdateTyConDef x y
        Case5 (x `PairE` LiftE y `PairE` z) -> UpdateFieldDef x y z
        _ -> error "impossible"
      _ -> error "impossible"
  
  instance SinkableE   TopEnvUpdate
  instance HoistableE  TopEnvUpdate
  instance RenameE     TopEnvUpdate
  
  instance GenericB TopEnvFrag where
    type RepB TopEnvFrag = PairB EnvFrag (LiftB (ModuleEnv `PairE` ListE TopEnvUpdate))
    fromB (TopEnvFrag x y (ReversedList z)) = PairB x (LiftB (y `PairE` ListE z))
    toB   (PairB x (LiftB (y `PairE` ListE z))) = TopEnvFrag x y (ReversedList z)
  
  instance RenameB     TopEnvFrag
  instance HoistableB  TopEnvFrag
  instance SinkableB TopEnvFrag
  instance ProvesExt   TopEnvFrag
  instance BindsNames  TopEnvFrag
  
  instance OutFrag TopEnvFrag where
    emptyOutFrag = TopEnvFrag emptyOutFrag mempty mempty
    {-# INLINE emptyOutFrag #-}
    catOutFrags (TopEnvFrag frag1 env1 partial1)
                (TopEnvFrag frag2 env2 partial2) =
      withExtEvidence frag2 $
        TopEnvFrag
          (catOutFrags frag1 frag2)
          (sink env1 <> env2)
          (sinkSnocList partial1 <> partial2)
    {-# INLINE catOutFrags #-}
  
  -- 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) mEnv' otherUpdates) = do
      let newerTopEnv = foldl applyUpdate newTopEnv otherUpdates
      Env newerTopEnv newModuleEnv
      where
        Env (TopEnv defs rules cache loadedM loadedO) mEnv = env
  
        newTopEnv = withExtEvidence frag $ TopEnv
          (defs `extendRecSubst` frag)
          (sink rules) (sink cache) (sink loadedM) (sink loadedO)
  
        newModuleEnv =
          ModuleEnv
            (imports <> imports')
            (sm   <> sm'   <> newImportedSM)
            (scs  <> scs'  <> newImportedSC)
          where
            ModuleEnv imports sm scs = withExtEvidence frag $ sink mEnv
            ModuleEnv imports' sm' scs' = 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
  
        lookupModulePure v = case lookupEnvPure newTopEnv v of ModuleBinding m -> m
  
  applyUpdate :: TopEnv n -> TopEnvUpdate n -> TopEnv n
  applyUpdate e = \case
    ExtendCache cache -> e { envCache = envCache e <> cache}
    AddCustomRule x y       -> e { envCustomRules   = envCustomRules   e <> CustomRules   (M.singleton x y)}
    UpdateLoadedModules x y -> e { envLoadedModules = envLoadedModules e <> LoadedModules (M.singleton x y)}
    UpdateLoadedObjects x y -> e { envLoadedObjects = envLoadedObjects e <> LoadedObjects (M.singleton x y)}
    FinishDictSpecialization dName methods -> do
      let SpecializedDictBinding (SpecializedDict dAbs oldMethods) = lookupEnvPure e dName
      case oldMethods of
        Nothing -> do
          let newBinding = SpecializedDictBinding $ SpecializedDict dAbs (Just methods)
          updateEnv dName newBinding e
        Just _ -> error "shouldn't be adding methods if we already have them"
    LowerDictSpecialization dName methods -> do
      let SpecializedDictBinding (SpecializedDict dAbs _) = lookupEnvPure e dName
      let newBinding = SpecializedDictBinding $ SpecializedDict dAbs (Just methods)
      updateEnv dName newBinding e
    UpdateTopFunEvalStatus f s -> do
      case lookupEnvPure e f of
        TopFunBinding (DexTopFun def lam _) ->
          updateEnv f (TopFunBinding $ DexTopFun def lam s) e
        _ -> error "can't update ffi function impl"
    UpdateInstanceDef name def -> do
      case lookupEnvPure e name of
        InstanceBinding _ ty -> updateEnv name (InstanceBinding def ty) e
    UpdateTyConDef name def -> do
      let TyConBinding _ methods = lookupEnvPure e name
      updateEnv name (TyConBinding (Just def) methods) e
    UpdateFieldDef name sn x -> do
      let TyConBinding def methods = lookupEnvPure e name
      updateEnv name (TyConBinding def (methods <> DotMethods (M.singleton sn x))) e
  
  updateEnv :: Color c => Name c n -> Binding c n -> TopEnv n -> TopEnv n
  updateEnv v rhs env =
    env { envDefs = RecSubst $ updateSubstFrag v rhs bs }
    where (RecSubst bs) = envDefs env
  
  lookupEnvPure :: Color c => TopEnv n -> Name c n -> Binding c n
  lookupEnvPure env v = lookupTerminalSubstFrag (fromRecSubst $ envDefs $ env) v
  
  instance GenericE Module where
    type RepE Module =       LiftE ModuleSourceName
                     `PairE` ListE ModuleName
                     `PairE` ListE ModuleName
                     `PairE` SourceMap
                     `PairE` SynthCandidates
  
    fromE (Module name deps transDeps sm sc) =
      LiftE name `PairE` ListE (S.toList deps) `PairE` ListE (S.toList transDeps)
        `PairE` sm `PairE` sc
    {-# INLINE fromE #-}
  
    toE (LiftE name `PairE` ListE deps `PairE` ListE transDeps
           `PairE` sm `PairE` sc) =
      Module name (S.fromList deps) (S.fromList transDeps) sm sc
    {-# INLINE toE #-}
  
  instance SinkableE      Module
  instance HoistableE     Module
  instance AlphaEqE       Module
  instance AlphaHashableE Module
  instance RenameE        Module
  
  instance GenericE ImportStatus where
    type RepE ImportStatus = ListE ModuleName `PairE` ListE ModuleName
    fromE (ImportStatus direct trans) = ListE (S.toList direct)
                                `PairE` ListE (S.toList trans)
    {-# INLINE fromE #-}
    toE (ListE direct `PairE` ListE trans) =
      ImportStatus (S.fromList direct) (S.fromList trans)
    {-# INLINE toE #-}
  
  instance SinkableE      ImportStatus
  instance HoistableE     ImportStatus
  instance AlphaEqE       ImportStatus
  instance AlphaHashableE ImportStatus
  instance RenameE        ImportStatus
  
  instance Semigroup (ImportStatus n) where
    ImportStatus direct trans <> ImportStatus direct' trans' =
      ImportStatus (direct <> direct') (trans <> trans')
  
  instance Monoid (ImportStatus n) where
    mappend = (<>)
    mempty = ImportStatus mempty mempty
  
  instance GenericE LoadedModules where
    type RepE LoadedModules = ListE (PairE (LiftE ModuleSourceName) ModuleName)
    fromE (LoadedModules m) =
      ListE $ M.toList m <&> \(v,md) -> PairE (LiftE v) md
    {-# INLINE fromE #-}
    toE (ListE pairs) =
      LoadedModules $ M.fromList $ pairs <&> \(PairE (LiftE v) md) -> (v, md)
    {-# INLINE toE #-}
  
  instance SinkableE      LoadedModules
  instance HoistableE     LoadedModules
  instance AlphaEqE       LoadedModules
  instance AlphaHashableE LoadedModules
  instance RenameE        LoadedModules
  
  instance GenericE LoadedObjects where
    type RepE LoadedObjects = ListE (PairE FunObjCodeName (LiftE NativeFunction))
    fromE (LoadedObjects m) =
      ListE $ M.toList m <&> \(v,p) -> PairE v (LiftE p)
    {-# INLINE fromE #-}
    toE (ListE pairs) =
      LoadedObjects $ M.fromList $ pairs <&> \(PairE v (LiftE p)) -> (v, p)
    {-# INLINE toE #-}
  
  instance SinkableE      LoadedObjects
  instance HoistableE     LoadedObjects
  instance RenameE        LoadedObjects
  
  instance GenericE ModuleEnv where
    type RepE ModuleEnv = ImportStatus
                  `PairE` SourceMap
                  `PairE` SynthCandidates
    fromE (ModuleEnv imports sm sc) = imports `PairE` sm `PairE` sc
    {-# INLINE fromE #-}
    toE   (imports `PairE` sm `PairE` sc) = ModuleEnv imports sm sc
    {-# INLINE toE #-}
  
  instance SinkableE      ModuleEnv
  instance HoistableE     ModuleEnv
  instance AlphaEqE       ModuleEnv
  instance AlphaHashableE ModuleEnv
  instance RenameE        ModuleEnv
  
  instance Semigroup (ModuleEnv n) where
    ModuleEnv x1 x2 x3 <> ModuleEnv y1 y2 y3 =
      ModuleEnv (x1<>y1) (x2<>y2) (x3<>y3)
  
  instance Monoid (ModuleEnv n) where
    mempty = ModuleEnv mempty mempty mempty
  
  instance Semigroup (LoadedModules n) where
    LoadedModules m1 <> LoadedModules m2 = LoadedModules (m2 <> m1)
  
  instance Monoid (LoadedModules n) where
    mempty = LoadedModules mempty
  
  instance Semigroup (LoadedObjects n) where
    LoadedObjects m1 <> LoadedObjects m2 = LoadedObjects (m2 <> m1)
  
  instance Monoid (LoadedObjects n) where
    mempty = LoadedObjects mempty
  
  
  -- === instance ===
  
  prettyRecord :: [(String, Doc ann)] -> Doc ann
  prettyRecord xs = foldMap (\(name, val) -> pretty name <> indented val) xs
  
  instance Pretty (TopEnv n) where
    pretty (TopEnv defs rules cache _ _) =
      prettyRecord [ ("Defs"          , pretty defs)
                   , ("Rules"         , pretty rules)
                   , ("Cache"         , pretty cache) ]
  
  instance Pretty (CustomRules n) where
    pretty _ = "TODO: Rule printing"
  
  instance Pretty (ImportStatus n) where
    pretty imports = pretty $ S.toList $ directImports imports
  
  instance Pretty (ModuleEnv n) where
    pretty (ModuleEnv imports sm sc) =
      prettyRecord [ ("Imports"         , pretty imports)
                   , ("Source map"      , pretty sm)
                   , ("Synth candidates", pretty sc) ]
  
  instance Pretty (Env n) where
    pretty (Env env1 env2) =
      prettyRecord [ ("Top env"   , pretty env1)
                   , ("Module env", pretty env2)]
  
  instance Pretty (SolverBinding n) where
    pretty (InfVarBound  ty) = "Inference variable of type:"   <+> pretty ty
    pretty (SkolemBound  ty) = "Skolem variable of type:"      <+> pretty ty
    pretty (DictBound    ty) = "Dictionary variable of type:"  <+> pretty ty
  
  instance Pretty (Binding c n) where
    pretty b = case b of
      -- using `unsafeCoerceIRE` here because otherwise we don't have `IRRep`
      -- TODO: can we avoid printing needing IRRep? Presumably it's related to
      -- manipulating sets or something, which relies on Eq/Ord, which relies on renaming.
      AtomNameBinding   info -> "Atom name:" <+> pretty (unsafeCoerceIRE @CoreIR info)
      TyConBinding dataDef _ -> "Type constructor: " <+> pretty dataDef
      DataConBinding tyConName idx -> "Data constructor:" <+>
        pretty tyConName <+> "Constructor index:" <+> pretty idx
      ClassBinding    classDef -> pretty classDef
      InstanceBinding instanceDef _ -> pretty instanceDef
      MethodBinding className idx -> "Method" <+> pretty idx <+> "of" <+> pretty className
      TopFunBinding f -> pretty f
      FunObjCodeBinding _ -> "<object file>"
      ModuleBinding  _ -> "<module>"
      PtrBinding   _ _ -> "<ptr>"
      SpecializedDictBinding _ -> "<specialized-dict-binding>"
      ImpNameBinding ty -> "Imp name of type: " <+> pretty ty
  
  instance Pretty (Module n) where
    pretty m = prettyRecord
      [ ("moduleSourceName"     , pretty $ moduleSourceName m)
      , ("moduleDirectDeps"     , pretty $ S.toList $ moduleDirectDeps m)
      , ("moduleTransDeps"      , pretty $ S.toList $ moduleTransDeps m)
      , ("moduleExports"        , pretty $ moduleExports m)
      , ("moduleSynthCandidates", pretty $ moduleSynthCandidates m) ]
  
  instance Pretty a => Pretty (EvalStatus a) where
    pretty = \case
      Waiting    -> "<waiting>"
      Running    -> "<running>"
      Finished a -> pretty a
  
  instance Pretty (EnvFrag n l) where
    pretty (EnvFrag bindings) = pretty bindings
  
  instance Pretty (Cache n) where
    pretty (Cache _ _ _ _ _ _) = "<cache>" -- TODO
  
  instance Pretty (SynthCandidates n) where
    pretty scs = "instance dicts:" <+> pretty (M.toList $ instanceDicts scs)
  
  instance Pretty (LoadedModules n) where
    pretty _ = "<loaded modules>"
  
  instance Pretty (TopFunDef n) where
    pretty = \case
      Specialization       s -> pretty s
      LinearizationPrimal  _ -> "<linearization primal>"
      LinearizationTangent _ -> "<linearization tangent>"
  
  instance Pretty (TopFun n) where
    pretty = \case
      DexTopFun def lam lowering ->
        "Top-level Function"
           <> hardline <+> "definition:" <+> pretty def
           <> hardline <+> "lambda:" <+> pretty lam
           <> hardline <+> "lowering:" <+> pretty lowering
      FFITopFun f _ -> pretty f
  
  instance IRRep r => Pretty (TopLam r n) where
    pretty (TopLam _ _ lam) = pretty lam
  
  instance IRRep r => Pretty (AtomBinding r n) where
    pretty binding = case binding of
      LetBound    b -> pretty b
      MiscBound   t -> pretty t
      SolverBound b -> pretty b
      FFIFunBound s _ -> pretty s
      NoinlineFun ty _ -> "Top function with type: " <+> pretty ty
      TopDataBound (RepVal ty _) -> "Top data with type: " <+> pretty ty
  
  instance Pretty (SpecializationSpec n) where
    pretty (AppSpecialization f (Abs bs (ListE args))) =
      "Specialization" <+> pretty f <+> pretty bs <+> pretty args
  
  instance Hashable a => Hashable (EvalStatus a)
  
  instance Store (SolverBinding n)
  instance IRRep r => Store (AtomBinding r n)
  instance IRRep r => Store (TopLam r n)
  instance Store (SynthCandidates n)
  instance Store (Module n)
  instance Store (ImportStatus n)
  instance Store (TopFunLowerings n)
  instance Store a => Store (EvalStatus a)
  instance Store (TopFun n)
  instance Store (TopFunDef n)
  instance Color c => Store (Binding c n)
  instance Store (ModuleEnv n)
  instance Store (SerializedEnv n)
  instance Store (AtomRules n)
  instance Store (LinearizationSpec n)
  instance Store (SpecializedDictDef n)
  instance Store (SpecializationSpec n)
  