chore: clean up IR.Arg usage (#9557)

src/Lean/Compiler/IR/Basic.lean

@@ -468,9 +468,9 @@ def VarId.alphaEqv (ρ : IndexRenaming) (v₁ v₂ : VarId) : Bool :=
 instance : AlphaEqv VarId := ⟨VarId.alphaEqv⟩
 
 def Arg.alphaEqv (ρ : IndexRenaming) : Arg → Arg → Bool
-  | Arg.var v₁,     Arg.var v₂     => aeqv ρ v₁ v₂
-  | Arg.erased,     Arg.erased     => true
-  | _,              _              => false
+  | .var v₁,     .var v₂     => aeqv ρ v₁ v₂
+  | .erased,     .erased     => true
+  | _,           _           => false
 
 instance : AlphaEqv Arg := ⟨Arg.alphaEqv⟩
 

src/Lean/Compiler/IR/Borrow.lean

@@ -156,8 +156,8 @@ def ownVar (x : VarId) : M Unit := do
 
 def ownArg (x : Arg) : M Unit :=
   match x with
-  | Arg.var x => ownVar x
-  | _         => pure ()
+  | .var x => ownVar x
+  | .erased => pure ()
 
 def ownArgs (xs : Array Arg) : M Unit :=
   xs.forM ownArg
@@ -212,8 +212,8 @@ def ownParamsUsingArgs (xs : Array Arg) (ps : Array Param) : M Unit :=
     let x := xs[i]
     let p := ps[i]!
     match x with
-    | Arg.var x => if (← isOwned x) then ownVar p.x
-    | _         => pure ()
+    | .var x => if (← isOwned x) then ownVar p.x
+    | .erased => pure ()
 
 /-- Mark `xs[i]` as owned if it is one of the parameters `ps`.
    We use this action to mark function parameters that are being "packed" inside constructors.
@@ -229,8 +229,8 @@ def ownArgsIfParam (xs : Array Arg) : M Unit := do
   let ctx ← read
   xs.forM fun x => do
     match x with
-    | Arg.var x => if ctx.paramSet.contains x.idx then ownVar x
-    | _ => pure ()
+    | .var x => if ctx.paramSet.contains x.idx then ownVar x
+    | .erased => pure ()
 
 def collectExpr (z : VarId) : Expr → M Unit
   | Expr.reset _ x      => ownVar z *> ownVar x
@@ -249,7 +249,7 @@ def collectExpr (z : VarId) : Expr → M Unit
 def preserveTailCall (x : VarId) (v : Expr) (b : FnBody) : M Unit := do
   let ctx ← read
   match v, b with
-  | (Expr.fap g ys), (FnBody.ret (Arg.var z)) =>
+  | (Expr.fap g ys), (FnBody.ret (.var z)) =>
     -- NOTE: we currently support TCO for self-calls only
     if ctx.currFn == g && x == z then
       let ps ← getParamInfo (ParamMap.Key.decl g)

src/Lean/Compiler/IR/Boxing.lean

@@ -52,18 +52,18 @@ def mkBoxedVersionAux (decl : Decl) : N Decl := do
     let p := ps[i]!
     let q := qs[i]
     if !p.ty.isScalar then
-      pure (newVDecls, xs.push (Arg.var q.x))
+      pure (newVDecls, xs.push (.var q.x))
     else
       let x ← N.mkFresh
-      pure (newVDecls.push (FnBody.vdecl x p.ty (Expr.unbox q.x) default), xs.push (Arg.var x))
+      pure (newVDecls.push (FnBody.vdecl x p.ty (Expr.unbox q.x) default), xs.push (.var x))
   let r ← N.mkFresh
   let newVDecls := newVDecls.push (FnBody.vdecl r decl.resultType (Expr.fap decl.name xs) default)
   let body ← if !decl.resultType.isScalar then
-    pure <| reshape newVDecls (FnBody.ret (Arg.var r))
+    pure <| reshape newVDecls (FnBody.ret (.var r))
   else
     let newR ← N.mkFresh
     let newVDecls := newVDecls.push (FnBody.vdecl newR .tobject (Expr.box decl.resultType r) default)
-    pure <| reshape newVDecls (FnBody.ret (Arg.var newR))
+    pure <| reshape newVDecls (FnBody.ret (.var newR))
   return Decl.fdecl (mkBoxedName decl.name) qs decl.resultType.boxed body decl.getInfo
 
 def mkBoxedVersion (decl : Decl) : Decl :=
@@ -222,17 +222,17 @@ def castArgsIfNeededAux (xs : Array Arg) (typeFromIdx : Nat → IRType) : M (Arr
   for x in xs do
     let expected := typeFromIdx i
     match x with
-    | Arg.erased =>
+    | .erased =>
       xs' := xs'.push x
-    | Arg.var x =>
+    | .var x =>
       let xType ← getVarType x
       if eqvTypes xType expected then
-        xs' := xs'.push (Arg.var x)
+        xs' := xs'.push (.var x)
       else
         let y ← M.mkFresh
         let v ← mkCast x xType expected
         let b := FnBody.vdecl y expected v FnBody.nil
-        xs' := xs'.push (Arg.var y)
+        xs' := xs'.push (.var y)
         bs := bs.push b
     i := i + 1
   return (xs', bs)

src/Lean/Compiler/IR/EmitC.lean

@@ -542,7 +542,7 @@ def emitVDecl (z : VarId) (t : IRType) (v : Expr) : M Unit :=
 def isTailCall (x : VarId) (v : Expr) (b : FnBody) : M Bool := do
   let ctx ← read;
   match v, b with
-  | Expr.fap f _, FnBody.ret (Arg.var y) => return f == ctx.mainFn && x == y
+  | Expr.fap f _, FnBody.ret (.var y) => return f == ctx.mainFn && x == y
   | _, _ => pure false
 
 def paramEqArg (p : Param) (x : Arg) : Bool :=

src/Lean/Compiler/IR/EmitUtil.lean

@@ -17,7 +17,7 @@ namespace Lean.IR
 /-- Return true iff `b` is of the form `let x := g ys; ret x` -/
 def isTailCallTo (g : Name) (b : FnBody) : Bool :=
   match b with
-  | FnBody.vdecl x _ (Expr.fap f _) (FnBody.ret (Arg.var y)) => x == y && f == g
+  | FnBody.vdecl x _ (Expr.fap f _) (FnBody.ret (.var y)) => x == y && f == g
   | _  => false
 
 def usesModuleFrom (env : Environment) (modulePrefix : Name) : Bool :=

src/Lean/Compiler/IR/Format.lean

@@ -14,8 +14,8 @@ namespace Lean
 namespace IR
 
 private def formatArg : Arg → Format
-  | Arg.var id => format id
-  | Arg.erased => "◾"
+  | .var id => format id
+  | .erased => "◾"
 
 instance : ToFormat Arg := ⟨private_decl% formatArg⟩
 

src/Lean/Compiler/IR/RC.lean

@@ -88,8 +88,8 @@ private def isBorrowParamAux (x : VarId) (ys : Array Arg) (consumeParamPred : Na
   ys.size.any fun i _ =>
     let y := ys[i]
     match y with
-    | Arg.erased => false
-    | Arg.var y  => x == y && !consumeParamPred i
+    | .erased => false
+    | .var y  => x == y && !consumeParamPred i
 
 private def isBorrowParam (x : VarId) (ys : Array Arg) (ps : Array Param) : Bool :=
   isBorrowParamAux x ys fun i => ! ps[i]!.borrow
@@ -103,15 +103,15 @@ private def getNumConsumptions (x : VarId) (ys : Array Arg) (consumeParamPred :
   ys.size.fold (init := 0) fun i _ n =>
     let y := ys[i]
     match y with
-    | Arg.erased => n
-    | Arg.var y  => if x == y && consumeParamPred i then n+1 else n
+    | .erased => n
+    | .var y  => if x == y && consumeParamPred i then n+1 else n
 
 private def addIncBeforeAux (ctx : Context) (xs : Array Arg) (consumeParamPred : Nat → Bool) (b : FnBody) (liveVarsAfter : LiveVarSet) : FnBody :=
   xs.size.fold (init := b) fun i _ b =>
     let x := xs[i]
     match x with
-    | Arg.erased => b
-    | Arg.var x =>
+    | .erased => b
+    | .var x =>
       let info := getVarInfo ctx x
       if !info.type.isPossibleRef || !isFirstOcc xs i then b
       else
@@ -131,8 +131,8 @@ private def addIncBefore (ctx : Context) (xs : Array Arg) (ps : Array Param) (b
 private def addDecAfterFullApp (ctx : Context) (xs : Array Arg) (ps : Array Param) (b : FnBody) (bLiveVars : LiveVarSet) : FnBody :=
 xs.size.fold (init := b) fun i _ b =>
   match xs[i] with
-  | Arg.erased => b
-  | Arg.var x  =>
+  | .erased => b
+  | .var x  =>
     /- We must add a `dec` if `x` must be consumed, it is alive after the application,
        and it has been borrowed by the application.
        Remark: `x` may occur multiple times in the application (e.g., `f x y x`).
@@ -203,7 +203,7 @@ private def processVDecl (ctx : Context) (z : VarId) (t : IRType) (v : Expr) (b
       addIncBefore ctx ys ps b bLiveVars
     | (Expr.pap _ ys)        => addIncBeforeConsumeAll ctx ys (FnBody.vdecl z t v b) bLiveVars
     | (Expr.ap x ys)         =>
-      let ysx := ys.push (Arg.var x) -- TODO: avoid temporary array allocation
+      let ysx := ys.push (.var x) -- TODO: avoid temporary array allocation
       addIncBeforeConsumeAll ctx ysx (FnBody.vdecl z t v b) bLiveVars
     | (Expr.unbox x)         => FnBody.vdecl z t v (addDecIfNeeded ctx x b bLiveVars)
     | _                      => FnBody.vdecl z t v b  -- Expr.reset, Expr.box, Expr.lit are handled here
@@ -259,10 +259,10 @@ partial def visitFnBody : FnBody → Context → (FnBody × LiveVarSet)
     (FnBody.case tid x xType alts, caseLiveVars)
   | b@(FnBody.ret x), ctx =>
     match x with
-    | Arg.var x =>
+    | .var x =>
       let info := getVarInfo ctx x
       if info.type.isPossibleRef && !info.consume then (addInc ctx x b, mkLiveVarSet x) else (b, mkLiveVarSet x)
-    | _         => (b, {})
+    | .erased => (b, {})
   | b@(FnBody.jmp j xs), ctx =>
     let jLiveVars := getJPLiveVars ctx j
     let ps        := getJPParams ctx j