Merge pull request #1318 from google-research/delivering-even-more-decls

More prep for decls-in-binders (again)

Author: dougalm

src/lib/Builder.hs

@@ -626,31 +626,14 @@ buildAbs hint binding cont = do
         return $ Abs b body
 {-# INLINE buildAbs #-}
 
-varsAsBinderNest :: (EnvReader m, IRRep r) => [AtomVar r n] -> m n (EmptyAbs (Nest (Binder r)) n)
-varsAsBinderNest [] = return $ EmptyAbs Empty
-varsAsBinderNest (v:vs) = do
-  rest <- varsAsBinderNest vs
-  ty <- return $ getType v
-  let AtomVar v' _ = v
-  Abs b (Abs bs UnitE) <- return $ abstractFreeVar v' rest
-  return $ EmptyAbs (Nest (b:>ty) bs)
-
 typesFromNonDepBinderNest
   :: (EnvReader m, Fallible1 m, IRRep r)
   => Nest (Binder r) n l -> m n [Type r n]
 typesFromNonDepBinderNest Empty = return []
-typesFromNonDepBinderNest (Nest (b:>ty) rest) = do
-  Abs rest' UnitE <- return $ ignoreHoistFailure $ hoist b (Abs rest UnitE)
+typesFromNonDepBinderNest (Nest b rest) = do
+  Abs rest' UnitE <- return $ assumeConst $ Abs (UnaryNest b) $ Abs rest UnitE
   tys <- typesFromNonDepBinderNest rest'
-  return $ ty : tys
-
-singletonBinderNest
-  :: (EnvReader m, IRRep r)
-  => NameHint -> ann n
-  -> m n (EmptyAbs (Nest (BinderP (AtomNameC r) ann)) n)
-singletonBinderNest hint ann = do
-  Abs b _ <- return $ newName hint
-  return $ EmptyAbs (Nest (b:>ann) Empty)
+  return $ binderType b : tys
 
 buildUnaryLamExpr
   :: (ScopableBuilder r m)
@@ -858,7 +841,7 @@ zeroAt ty = liftEmitBuilder $ go ty where
    BaseTy bt  -> return $ Con $ Lit $ zeroLit bt
    ProdTy tys -> ProdVal <$> mapM go tys
    TabPi tabPi -> buildFor (getNameHint tabPi) Fwd (tabIxType tabPi) \i ->
-     go =<< instantiateTabPiTy (sink tabPi) (Var i)
+     go =<< instantiate (sink tabPi) [Var i]
    _ -> unreachable
   zeroLit bt = case bt of
     Scalar Float64Type -> Float64Lit 0.0
@@ -1134,8 +1117,7 @@ naryTopAppInlined :: (Builder SimpIR m, Emits n) => TopFunName n -> [SAtom n] ->
 naryTopAppInlined f xs = do
   TopFunBinding f' <- lookupEnv f
   case f' of
-    DexTopFun _ (TopLam _ _ (LamExpr bs body)) _ ->
-      applySubst (bs@@>(SubstVal<$>xs)) body >>= emitBlock
+    DexTopFun _ lam _ -> instantiate lam xs >>= emitBlock
     _ -> naryTopApp f xs
 {-# INLINE naryTopAppInlined #-}
 
@@ -1194,8 +1176,7 @@ applyIxMethod dict method args = case dict of
   IxDictSpecialized _ d params -> do
     SpecializedDict _ maybeFs <- lookupSpecDict d
     Just fs <- return maybeFs
-    TopLam _ _ (LamExpr bs body) <- return $ fs !! fromEnum method
-    emitBlock =<< applySubst (bs @@> fmap SubstVal (params ++ args)) body
+    instantiate (fs !! fromEnum method) (params ++ args) >>= emitBlock
 
 unsafeFromOrdinal :: (SBuilder m, Emits n) => IxType SimpIR n -> Atom SimpIR n -> m n (Atom SimpIR n)
 unsafeFromOrdinal (IxType _ dict) i = applyIxMethod dict UnsafeFromOrdinal [i]
@@ -1262,10 +1243,10 @@ isJustE x = liftEmitBuilder $
 -- Monoid a -> (n=>a) -> a
 reduceE :: (Emits n, Builder r m) => BaseMonoid r n -> Atom r n -> m n (Atom r n)
 reduceE monoid xs = liftEmitBuilder do
-  TabTy d (n:>ty) a <- return $ getType xs
-  a' <- return $ ignoreHoistFailure $ hoist n a
-  getSnd =<< emitRunWriter noHint a' monoid \_ ref ->
-    buildFor noHint Fwd (sink $ IxType ty d) \i -> do
+  TabPi tabPi <- return $ getType xs
+  let a = assumeConst tabPi
+  getSnd =<< emitRunWriter noHint a monoid \_ ref ->
+    buildFor noHint Fwd (sink $ tabIxType tabPi) \i -> do
       x <- tabApp (sink xs) (Var i)
       emitExpr $ PrimOp $ RefOp (sink $ Var ref) $ MExtend (sink monoid) x
 
@@ -1278,11 +1259,10 @@ andMonoid = liftM (BaseMonoid TrueAtom) $ liftBuilder $
 mapE :: (Emits n, ScopableBuilder r m)
      => (forall l. (Emits l, DExt n l) => Atom r l -> m l (Atom r l))
      -> Atom r n -> m n (Atom r n)
-mapE f xs = do
-  TabTy d (n:>ty) _ <- return $ getType xs
-  buildFor (getNameHint n) Fwd (IxType ty d) \i -> do
-    x <- tabApp (sink xs) (Var i)
-    f x
+mapE cont xs = do
+  TabPi tabPi <- return $ getType xs
+  buildFor (getNameHint tabPi) Fwd (tabIxType tabPi) \i -> do
+    tabApp (sink xs) (Var i) >>= cont
 
 -- (n:Type) ?-> (a:Type) ?-> (xs : n=>Maybe a) : Maybe (n => a) =
 catMaybesE :: (Emits n, Builder r m) => Atom r n -> m n (Atom r n)

src/lib/CheapReduction.hs

@@ -11,12 +11,12 @@ module CheapReduction
   ( CheaplyReducibleE (..), cheapReduce, cheapReduceWithDecls, cheapNormalize
   , normalizeProj, asNaryProj, normalizeNaryProj
   , depPairLeftTy, instantiateTyConDef
-  , dataDefRep, instantiateDepPairTy, unwrapNewtypeType, repValAtom
+  , dataDefRep, unwrapNewtypeType, repValAtom
   , unwrapLeadingNewtypesType, wrapNewtypesData, liftSimpAtom, liftSimpType
   , liftSimpFun, makeStructRepVal, NonAtomRenamer (..), Visitor (..), VisitGeneric (..)
   , visitAtomPartial, visitTypePartial, visitAtomDefault, visitTypeDefault, Visitor2
-  , visitBinders, visitPiDefault, visitAlt, toAtomVar, instantiatePiTy, instantiateTabPiTy
-  , bindersToVars, bindersToAtoms)
+  , visitBinders, visitPiDefault, visitAlt, toAtomVar, instantiate, withInstantiated
+  , bindersToVars, bindersToAtoms, instantiateNames, withInstantiatedNames, assumeConst)
   where
 
 import Control.Applicative
@@ -242,7 +242,7 @@ cheapReduceDictExpr resultTy d = case d of
         args' <- mapM cheapReduceE args
         InstanceDef _ _ bs _ body <- lookupInstanceDef instanceName
         let InstanceBody superclasses _ = body
-        applySubst (bs@@>(SubstVal <$> args')) (superclasses !! superclassIx)
+        instantiate (Abs bs (superclasses !! superclassIx)) args'
       child' -> return $ DictCon resultTy $ SuperclassProj child' superclassIx
   InstantiatedGiven f xs ->
     reduceApp <|> justSubst
@@ -285,19 +285,16 @@ instance IRRep r => CheaplyReducibleE r (Expr r) (Atom r) where
       cheapReduceE dict >>= \case
         DictCon _ (InstanceDict instanceName args) -> dropSubst do
           args' <- mapM cheapReduceE args
-          InstanceDef _ _ bs _ (InstanceBody _ methods) <- lookupInstanceDef instanceName
-          let method = methods !! i
-          extendSubst (bs@@>(SubstVal <$> args')) do
-            method' <- cheapReduceE method
+          def <- lookupInstanceDef instanceName
+          withInstantiated def args' \(PairE _ (InstanceBody _ methods)) -> do
+            method' <- cheapReduceE $ methods !! i
             cheapReduceApp method' explicitArgs'
         _ -> empty
     _ -> empty
 
 cheapReduceApp :: CAtom o -> [CAtom o] -> CheapReducerM CoreIR i o (CAtom o)
 cheapReduceApp f xs = case f of
-  Lam (CoreLamExpr _ (LamExpr bs body)) -> do
-    let subst = bs @@> fmap SubstVal xs
-    dropSubst $ extendSubst subst $ cheapReduceE body
+  Lam lam -> dropSubst $ withInstantiated lam xs \body -> cheapReduceE body
   _ -> empty
 
 instance IRRep r => CheaplyReducibleE r (IxType r) (IxType r) where
@@ -450,7 +447,7 @@ projType i ty x = case ty of
   DepPairTy t | i == 0 -> return $ depPairLeftTy t
   DepPairTy t | i == 1 -> do
     xFst <- normalizeProj (ProjectProduct 0) x
-    instantiateDepPairTy t xFst
+    instantiate t [xFst]
   _ -> error $ "Can't project type: " ++ pprint ty
 
 unwrapLeadingNewtypesType :: EnvReader m => CType n -> m n ([NewtypeCon n], CType n)
@@ -470,13 +467,39 @@ instantiateTyConDef (TyConDef _ _ bs conDefs) (TyConParams _ xs) = do
   applySubst (bs @@> (SubstVal <$> xs)) conDefs
 {-# INLINE instantiateTyConDef #-}
 
-instantiatePiTy :: (EnvReader m, IRRep r) => PiType r n -> [Atom r n] -> m n (EffTy r n)
-instantiatePiTy (PiType bs effTy) xs = do
-  applySubst (bs @@> (SubstVal <$> xs)) effTy
-
-instantiateTabPiTy :: (EnvReader m, IRRep r) => TabPiType r n -> Atom r n -> m n (Type r n)
-instantiateTabPiTy (TabPiType _ b resultTy) x = do
-  applySubst (b @> SubstVal x) resultTy
+assumeConst
+  :: (IRRep r, HoistableE body, SinkableE body, ToBindersAbs e body r) => e n -> body n
+assumeConst e = case toAbs e of Abs bs body -> ignoreHoistFailure $ hoist bs body
+
+instantiate
+  :: (EnvReader m, IRRep r, SubstE (SubstVal Atom) body, SinkableE body, ToBindersAbs e body r)
+  => e n -> [Atom r n] -> m n (body n)
+instantiate e xs = case toAbs e of
+  Abs bs body -> applySubst (bs @@> (SubstVal <$> xs)) body
+
+-- "lazy" subst-extending version of `instantiate`
+withInstantiated
+  :: (SubstReader AtomSubstVal m, IRRep r, SubstE (SubstVal Atom) body, SinkableE body, ToBindersAbs e body r)
+  => e i -> [Atom r o]
+  -> (forall i'. body i' -> m i' o a)
+  -> m i o a
+withInstantiated e xs cont = case toAbs e of
+  Abs bs body -> extendSubst (bs @@> (SubstVal <$> xs)) $ cont body
+
+instantiateNames
+  :: (EnvReader m, IRRep r, RenameE body, SinkableE body, ToBindersAbs e body r)
+  => e n -> [AtomName r n] -> m n (body n)
+instantiateNames e vs = case toAbs e of
+  Abs bs body -> applyRename (bs @@> vs) body
+
+-- "lazy" subst-extending version of `instantiateNames`
+withInstantiatedNames
+  :: (SubstReader Name m, IRRep r, RenameE body, SinkableE body, ToBindersAbs e body r)
+  => e i -> [AtomName r o]
+  -> (forall i'. body i' -> m i' o a)
+  -> m i o a
+withInstantiatedNames e vs cont = case toAbs e of
+  Abs bs body -> extendRenamer (bs @@> vs) $ cont body
 
 -- Returns a representation type (type of an TypeCon-typed Newtype payload)
 -- given a list of instantiated DataConDefs.
@@ -498,10 +521,6 @@ makeStructRepVal tyConName args = do
       _ -> error "wrong number of args"
     _ -> return $ ProdVal args
 
-instantiateDepPairTy :: (IRRep r, EnvReader m) => DepPairType r n -> Atom r n -> m n (Type r n)
-instantiateDepPairTy (DepPairType _ b rhsTy) x = applyAbs (Abs b rhsTy) (SubstVal x)
-{-# INLINE instantiateDepPairTy #-}
-
 -- === traversable terms ===
 
 class Monad m => NonAtomRenamer m i o | m -> i, m -> o where

src/lib/CheckType.hs

@@ -214,7 +214,7 @@ instance IRRep r => HasType r (Atom r) where
     DepPair l r ty -> do
       ty' <- checkTypeE TyKind ty
       l'  <- checkTypeE (depPairLeftTy ty') l
-      rTy <- instantiateDepPairTy ty' l'
+      rTy <- instantiate ty' [l']
       r |: rTy
       return $ DepPairTy ty'
     Con con  -> typeCheckPrimCon con
@@ -236,7 +236,7 @@ instance IRRep r => HasType r (Atom r) where
         DepPairTy t | i == 1 -> do
           x' <- renameM x
           xFst <- normalizeProj (ProjectProduct 0) x'
-          instantiateDepPairTy t xFst
+          instantiate t [xFst]
         _ -> throw TypeErr $ "Not a product type:" ++ pprint ty
     TypeAsAtom ty -> getTypeE ty
 
@@ -275,7 +275,7 @@ instance IRRep r => HasType r (Type r) where
         DepPairTy t | i == 1 -> do
           x' <- renameM x
           xFst <- normalizeProj (ProjectProduct 0) x'
-          instantiateDepPairTy t xFst
+          instantiate t [xFst]
         _ -> throw TypeErr $ "Not a product type:" ++ pprint ty
 
 instance HasType CoreIR SimpInCore where

src/lib/Imp.hs

@@ -373,7 +373,7 @@ toImpRefOp refDest' m = do
               idx <- unsafeFromOrdinalImp (sink ixTy) i
               xElt <- liftBuilderImp $ tabApp (sink x) (sink idx)
               yElt <- liftBuilderImp $ tabApp (sink y) (sink idx)
-              eltTy <- instantiateTabPiTy (sink t) idx
+              eltTy <- instantiate (sink t) [idx]
               ithDest <- indexDest (sink accDest) idx
               liftMonoidCombine ithDest eltTy (sink bc) xElt yElt
           _ -> error $ "Base monoid type mismatch: can't lift " ++
@@ -584,7 +584,7 @@ toImpTypedHof (TypedHof (EffTy _ resultTy') hof) = do
               emitLoop noHint Fwd n \i -> do
                 idx <- unsafeFromOrdinalImp (sink ixTy) i
                 x' <- sinkM x
-                eltTy <- instantiateTabPiTy (sink t) idx
+                eltTy <- instantiate (sink t) [idx]
                 ithDest <- indexDest (sink accDest) idx
                 liftMonoidEmpty ithDest eltTy x'
             _ -> error $ "Base monoid type mismatch: can't lift " ++
@@ -1003,7 +1003,7 @@ buildGarbageVal ty =
 
 indexDest :: Emits n => Dest n -> SAtom n -> SubstImpM i n (Dest n)
 indexDest (Dest (TabPi tabTy) tree) i = do
-  eltTy <- instantiateTabPiTy tabTy i
+  eltTy <- instantiate tabTy [i]
   ord <- ordinalImp (tabIxType tabTy) i
   leafTys <- typeToTree $ TabPi tabTy
   Dest eltTy <$> forM (zipTrees leafTys tree) \(leafTy, ptr) -> do
@@ -1027,7 +1027,7 @@ indexRepValParam :: Emits n
   -> (IExpr n -> SubstImpM i n (IExpr n))
   -> SubstImpM i n (SRepVal n)
 indexRepValParam (RepVal (TabPi tabTy) vals) i tyFunc func = do
-  eltTy <- instantiateTabPiTy tabTy i
+  eltTy <- instantiate tabTy [i]
   ord <- ordinalImp (tabIxType tabTy) i
   leafTys <- typeToTree (TabPi tabTy)
   vals' <- forM (zipTrees leafTys vals) \(leafTy, ptr) -> do

src/lib/Inference.hs

@@ -887,7 +887,7 @@ checkSigma hint expr sTy = confuseGHC >>= \_ -> case sTy of
       -- TODO: check for the case that we're given some of the implicit dependent pair args explicitly
       lhsVal <- Var <$> freshInferenceName MiscInfVar lhsTy
       -- TODO: make an InfVarDesc case for dep pair instantiation
-      rhsTy <- instantiateDepPairTy depPairTy lhsVal
+      rhsTy <- instantiate depPairTy [lhsVal]
       rhsVal <- checkSigma noHint expr rhsTy
       return $ DepPair lhsVal rhsVal depPairTy
     _ -> fallback
@@ -996,7 +996,7 @@ checkOrInferRho hint uExprWithSrc@(WithSrcE pos expr) reqTy = do
     case reqTy of
       Check (DepPairTy ty@(DepPairType _ (_ :> lhsTy) _)) -> do
         lhs' <- checkSigmaDependent noHint lhs lhsTy
-        rhsTy <- instantiateDepPairTy ty lhs'
+        rhsTy <- instantiate ty [lhs']
         rhs' <- checkSigma noHint rhs rhsTy
         return $ DepPair lhs' rhs' ty
       _ -> throw TypeErr $ "Can't infer the type of a dependent pair; please annotate it"

src/lib/Lower.hs

@@ -67,7 +67,7 @@ lowerFullySequential wantDestStyle (TopLam False piTy (LamExpr bs body)) = liftE
     True -> do
       refreshAbs (Abs bs body) \bs' body' -> do
         xs <- bindersToAtoms bs'
-        EffTy _ resultTy <- instantiatePiTy (sink piTy) xs
+        EffTy _ resultTy <- instantiate (sink piTy) xs
         Abs b body'' <- lowerFullySequentialBlock resultTy body'
         return $ LamExpr (bs' >>> UnaryNest b) body''
     False -> do