feat(library/init): add more general modify: modifyGet

This commit also adds a few helper elaboration functions.

library/init/control/estate.lean

@@ -72,9 +72,11 @@ fun r => match r with
   | ⟨Result.ok _ s, _⟩    => Result.ok s s
   | ⟨Result.error _ _, h⟩ => unreachableError h
 
-@[inline] protected def modify (f : σ → σ) : EState ε σ PUnit :=
+@[inline] protected def modifyGet (f : σ → α × σ) : EState ε σ α :=
 fun r => match r with
-  | ⟨Result.ok _ s, _⟩    => Result.ok ⟨⟩ (f s)
+  | ⟨Result.ok _ s, _⟩    =>
+    match f s with
+    | (a, s) => Result.ok a s
   | ⟨Result.error _ _, h⟩ => unreachableError h
 
 @[inline] protected def throw (e : ε) : EState ε σ α :=
@@ -129,7 +131,7 @@ instance : HasOrelse (EState ε σ α) :=
 { orelse := @EState.orelse _ _ _ }
 
 instance : MonadState σ (EState ε σ) :=
-{ set := @EState.set _ _, get := @EState.get _ _, modify := @EState.modify _ _ }
+{ set := @EState.set _ _, get := @EState.get _ _, modifyGet := @EState.modifyGet _ _ }
 
 instance : MonadExcept ε (EState ε σ) :=
 { throw := @EState.throw _ _, catch := @EState.catch _ _ }

library/init/control/state.lean

@@ -55,8 +55,8 @@ fun s => pure (s, s)
 @[inline] protected def set : σ → StateT σ m PUnit :=
 fun s' s => pure (⟨⟩, s')
 
-@[inline] protected def modify (f : σ → σ) : StateT σ m PUnit :=
-fun s => pure (⟨⟩, f s)
+@[inline] protected def modifyGet (f : σ → α × σ) : StateT σ m α :=
+fun s => pure (f s)
 
 @[inline] protected def lift {α : Type u} (t : m α) : StateT σ m α :=
 fun s => do a ← t; pure (a, s)
@@ -90,15 +90,18 @@ class MonadState (σ : outParam (Type u)) (m : Type u → Type v) :=
 (set {} : σ → m PUnit)
 /- Map the top-most State of a Monad stack.
 
-   Note: `modify f` may be preferable to `f <$> get >>= put` because the latter
-   does not use the State linearly (without sufficient inlining). -/
-(modify {} : (σ → σ) → m PUnit)
+   Note: `modifyGet f` may be preferable to `do s <- get; let (a, s) := f s; put s; pure a`
+   because the latter does not use the State linearly (without sufficient inlining). -/
+(modifyGet {} {α : Type u} : (σ → α × σ) → m α)
 
-export MonadState (get set modify)
+export MonadState (get set modifyGet)
 
 section
 variables {σ : Type u} {m : Type u → Type v}
 
+@[inline] def modify [MonadState σ m] (f : σ → σ) : m PUnit :=
+modifyGet (fun s => (PUnit.unit, f s))
+
 @[inline] def getModify [MonadState σ m] [Monad m] (f : σ → σ) : m σ :=
 do s ← get; modify f; pure s
 
@@ -107,12 +110,13 @@ do s ← get; modify f; pure s
 instance monadStateTrans {n : Type u → Type w} [HasMonadLift m n] [MonadState σ m] : MonadState σ n :=
 { get := monadLift (MonadState.get : m _),
   set := fun st => monadLift (MonadState.set st : m _),
-  modify := fun f => monadLift (MonadState.modify f : m _) }
+  modifyGet := fun α f => monadLift (MonadState.modifyGet f : m _) }
 
 instance [Monad m] : MonadState σ (StateT σ m) :=
 { get := StateT.get,
   set := StateT.set,
-  modify := StateT.modify }
+  modifyGet := @StateT.modifyGet _ _ _ }
+
 end
 
 /-- Adapt a Monad stack, changing the Type of its top-most State.

library/init/lean/compiler/ir/boxing.lean

@@ -43,9 +43,7 @@ def isBoxedName : Name → Bool
 abbrev N := State Nat
 
 private def mkFresh : N VarId :=
-do idx ← get;
-   modify (fun n => n + 1);
-   pure {idx := idx }
+modifyGet $ fun n => ({ idx := n }, n + 1)
 
 def requiresBoxedVersion (env : Environment) (decl : Decl) : Bool :=
 let ps := decl.params;

library/init/lean/compiler/ir/emitutil.lean

@@ -55,13 +55,12 @@ def usesLeanNamespace (env : Environment) : Decl → Bool
 | Decl.fdecl _ _ _ b => (UsesLeanNamespace.visitFnBody b env).run' {}
 | _                  => false
 
-
 namespace CollectUsedDecls
 
 abbrev M := ReaderT Environment (State NameSet)
 
 @[inline] def collect (f : FunId) : M Unit :=
-modify (fun s => s.insert f)
+modify $ fun s => s.insert f
 
 partial def collectFnBody : FnBody → M Unit
 | FnBody.vdecl _ _ v b   =>

library/init/lean/compiler/ir/expandresetreuse.lean

@@ -137,7 +137,7 @@ mask.foldl
 
 abbrev M := ReaderT Context (State Nat)
 def mkFresh : M VarId :=
-do idx ← get; modify (fun n => n + 1); pure { idx := idx }
+modifyGet $ fun n => ({ idx := n }, n + 1)
 
 def releaseUnreadFields (y : VarId) (mask : Mask) (b : FnBody) : M FnBody :=
 mask.size.mfold

library/init/lean/compiler/ir/normids.lean

@@ -15,9 +15,9 @@ namespace UniqueIds
 abbrev M := StateT IndexSet Id
 
 def checkId (id : Index) : M Bool :=
-do found ← get;
-   if found.contains id then pure false
-   else modify (fun s => s.insert id) *> pure true
+modifyGet $ fun s =>
+  if s.contains id then (false, s)
+  else (true, s.insert id)
 
 def checkParams (ps : Array Param) : M Bool :=
 ps.allM $ fun p => checkId p.x.idx

library/init/lean/elaborator/basic.lean

@@ -31,13 +31,14 @@ structure ElabContext :=
 (fileMap  : FileMap)
 
 structure ElabScope :=
-(cmd       : String)
-(header    : Name)
-(options   : Options := {})
-(ns        : Name := Name.anonymous) -- current namespace
-(openDecls : List OpenDecl := [])
-(univs     : List Name := [])
-(lctx      : LocalContext := {})
+(cmd         : String)
+(header      : Name)
+(options     : Options := {})
+(ns          : Name := Name.anonymous) -- current namespace
+(openDecls   : List OpenDecl := [])
+(univs       : List Name := [])
+(lctx        : LocalContext := {})
+(nextInstIdx : Nat := 1)
 
 namespace ElabScope
 
@@ -399,10 +400,7 @@ instance {α} : Inhabited (Elab α) :=
 ⟨fun _ => default _⟩
 
 def mkFreshName : Elab Name :=
-do s ← get;
-   let n := s.ngen.curr;
-   modify $ fun s => { ngen := s.ngen.next, .. s };
-   pure n
+modifyGet $ fun s => (s.ngen.curr, { ngen := s.ngen.next, .. s })
 
 def getScope : Elab ElabScope :=
 do s ← get; pure s.scopes.head
@@ -419,6 +417,37 @@ do s ← get;
    | []      => pure Name.anonymous
    | (sc::_) => pure sc.ns
 
+@[specialize] def modifyScope (f : ElabScope → ElabScope) : Elab Unit :=
+modify $ fun s =>
+  { scopes := match s.scopes with
+    | h::t => f h :: t
+    | []   => [], -- unreachable
+    .. s }
+
+@[specialize] def modifyGetScope {α} [Inhabited α] (f : ElabScope → α × ElabScope) : Elab α :=
+modifyGet $ fun s =>
+  match s with
+  | { scopes := h::t, .. } =>
+    let (a, h) := f h;
+    (a, { scopes := h :: t, .. s })
+  | _ => (default _, s)
+
+def mkLocalDecl (userName : Name) (type : Expr) (bi : BinderInfo := BinderInfo.default) : Elab LocalDecl :=
+do
+idx ← mkFreshName;
+modifyGetScope $ fun scope =>
+   let (decl, lctx) := scope.lctx.mkLocalDecl idx userName type bi;
+   (decl, { lctx := lctx, .. scope })
+
+def anonymousInstNamePrefix := `_inst
+
+def mkAnonymousInstName : Elab Name :=
+do
+scope ← getScope;
+let n := anonymousInstNamePrefix.appendIndexAfter scope.nextInstIdx;
+modifyScope $ fun scope => { nextInstIdx := scope.nextInstIdx + 1, .. scope };
+pure n
+
 def rootNamespace := `_root_
 
 def removeRoot (n : Name) : Name :=
@@ -453,6 +482,13 @@ do s ← get;
        | some n => pure n
        | none   => throw (ElabException.other ("unknown namespace '" ++ toString n ++ "'"))
 
+@[inline] def withNewScope {α} (x : Elab α) : Elab α :=
+do
+modify $ fun s => { scopes := s.scopes.head :: s.scopes, .. s };
+a ← x;
+modify $ fun s => { scopes := s.scopes.tail, .. s};
+pure a
+
 /- Remark: in an ideal world where performance doesn't matter, we would define `Elab` as
    ```
    ExceptT ElabException (StateT ElabException IO)

library/init/lean/elaborator/command.lean

@@ -88,13 +88,6 @@ fun n => do
     [];
   modify $ fun s => { env := aliases.foldl (fun env p => addAlias env p.1 p.2) s.env, .. s }
 
-@[specialize] def modifyScope (f : ElabScope → ElabScope) : Elab Unit :=
-modify $ fun s =>
-  { scopes := match s.scopes with
-    | h::t => f h :: t
-    | []   => [], -- unreachable
-    .. s }
-
 def addOpenDecl (d : OpenDecl) : Elab Unit :=
 modifyScope $ fun scope => { openDecls := d :: scope.openDecls, .. scope }
 

library/init/lean/elaborator/preterm.lean

@@ -114,6 +114,9 @@ stx.ifNode
     | none    => logErrorAndThrow stx ("`toPreTerm` failed, no support for syntax '" ++ toString k ++ "'"))
   (fun _ => throw "`toPreTerm` failed, unexpected syntax")
 
+private def mkHoleFor (stx : Syntax Expr) : Elab PreTerm :=
+setPos stx (Expr.mvar Name.anonymous)
+
 @[builtinPreTermElab «type»] def convertType : PreTermElab :=
 fun _ => pure $ Expr.sort $ Level.succ Level.zero
 
@@ -132,6 +135,42 @@ fun n => do
    else
      pure $ Expr.sort level
 
+private def mkLocal (decl : LocalDecl) : PreTerm :=
+Expr.local decl.name decl.userName decl.type decl.binderInfo
+
+private def processBinder (b : Syntax Expr) : Elab (List PreTerm) :=
+match b.getKind with
+| `Lean.Parser.Term.simpleBinder   => do
+   let args := (b.getArg 0).getArgs;
+   args.mfoldl (fun r arg => do
+     let id := arg.getId;
+     hole ← mkHoleFor arg;
+     decl ← mkLocalDecl id hole;
+     pure (mkLocal decl :: r))
+     []
+| `Lean.Parser.Term.explicitBinder => do runIO (IO.println $ ">> explicit " ++ (toString b)); pure []
+| `Lean.Parser.Term.implicitBinder => do runIO (IO.println $ ">> implict " ++ (toString b)); pure []
+| `Lean.Parser.Term.instBinder     => do runIO (IO.println $ ">> inst " ++ (toString b)); pure []
+| _ => throw "unknown binder kind"
+
+private def processBinders (bs : Array (Syntax Expr)) : Elab (List PreTerm) :=
+bs.mfoldl (fun r s => do xs ← processBinder s; pure (r ++ xs)) []
+
+@[builtinPreTermElab «forall»] def convertForall : PreTermElab :=
+fun n => do
+  let binders := n.getArg 1;
+  let body    := n.getArg 3;
+  withNewScope $ do
+    xs   ← processBinders binders.getArgs;
+    runIO (IO.println (xs.map Expr.dbgToString));
+    body ← toPreTerm body;
+    -- TODO
+    pure $ Expr.sort (Level.zero)
+--  dom ← toPreTerm $ n.getArg 0;
+--  rng ← toPreTerm $ n.getArg 2;
+--  pure $ Expr.pi id BinderInfo.default dom rng
+
+-- TODO: delete... arrow should be expanded at term.lean
 @[builtinPreTermElab «arrow»] def convertArrow : PreTermElab :=
 fun n => do
   id ← mkFreshName;

library/init/lean/elaborator/term.lean

@@ -53,5 +53,11 @@ fun stx _ => do
   let nilId    := closeBkt.mkIdent `List.nil;
   pure $ args.foldSepArgs (fun arg r => mkAppStx consId [arg, r]) nilId
 
+@[builtinTermElab arrow] def elabArrow : TermElab :=
+fun stx _ => do
+  id ← mkFreshName;
+  runIO (IO.println stx.val);
+  pure $ Syntax.other $ Expr.sort (Level.zero)
+
 end Elab
 end Lean