perf: use IR type info to decide whether to insert RC ops (#9396)

This is mostly a refactoring that replaces other analyses with type
information, but due to the introduction of `tagged` it also has the
side effect of eliminating ref counting ops entirely for types that
always have a tagged scalar representation, e.g. `Unit`.

src/Lean/Compiler/IR/Basic.lean

@@ -100,6 +100,10 @@ def isObj : IRType → Bool
   | tobject => true
   | _       => false
 
+def isPossibleRef : IRType → Bool
+  | object | tobject => true
+  | _ => false
+
 def isErased : IRType → Bool
   | erased => true
   | _ => false

src/Lean/Compiler/IR/RC.lean

@@ -16,7 +16,7 @@ that introduce the instructions `release` and `set`
 -/
 
 structure VarInfo where
-  ref        : Bool -- true if the variable may be a reference (aka pointer) at runtime
+  type       : IRType
   persistent : Bool -- true if the variable is statically known to be marked a Persistent at runtime
   consume    : Bool -- true if the variable RC must be "consumed"
   deriving Inhabited
@@ -52,7 +52,7 @@ def getJPLiveVars (ctx : Context) (j : JoinPointId) : LiveVarSet :=
 
 def mustConsume (ctx : Context) (x : VarId) : Bool :=
   let info := getVarInfo ctx x
-  info.ref && info.consume
+  info.type.isPossibleRef && info.consume
 
 @[inline] def addInc (ctx : Context) (x : VarId) (b : FnBody) (n := 1) : FnBody :=
   let info := getVarInfo ctx x
@@ -69,7 +69,7 @@ private def updateRefUsingCtorInfo (ctx : Context) (x : VarId) (c : CtorInfo) :
     let m := ctx.varMap
     { ctx with
       varMap := match m.find? x with
-      | some info => m.insert x { info with ref := false } -- I really want a Lenses library + notation
+      | some info => m.insert x { info with type := c.type }
       | none      => m }
 
 private def addDecForAlt (ctx : Context) (caseLiveVars altLiveVars : LiveVarSet) (b : FnBody) : FnBody :=
@@ -112,7 +112,7 @@ private def addIncBeforeAux (ctx : Context) (xs : Array Arg) (consumeParamPred :
     | Arg.erased => b
     | Arg.var x =>
       let info := getVarInfo ctx x
-      if !info.ref || !isFirstOcc xs i then b
+      if !info.type.isPossibleRef || !isFirstOcc xs i then b
       else
         let numConsuptions := getNumConsumptions x xs consumeParamPred -- number of times the argument is
         let numIncs :=
@@ -162,16 +162,21 @@ private def consumeExpr (m : VarMap) : Expr → Bool
 
 /-- Return true iff `v` at runtime is a scalar value stored in a tagged pointer.
    We do not need RC operations for this kind of value. -/
-private def isScalarBoxedInTaggedPtr (v : Expr) : Bool :=
+private def typeForScalarBoxedInTaggedPtr? (v : Expr) : Option IRType :=
   match v with
-  | Expr.ctor c _           => c.size == 0 && c.ssize == 0 && c.usize == 0
-  | Expr.lit (LitVal.num n) => n ≤ maxSmallNat
-  | _ => false
+  | .ctor c _ =>
+    some c.type
+  | .lit (.num n) =>
+    if n ≤ maxSmallNat then
+      some .tagged
+    else
+      some .tobject
+  | _ => none
 
 private def updateVarInfo (ctx : Context) (x : VarId) (t : IRType) (v : Expr) : Context :=
   { ctx with
     varMap := ctx.varMap.insert x {
-        ref := t.isObj && !isScalarBoxedInTaggedPtr v,
+        type := typeForScalarBoxedInTaggedPtr? v |>.getD t
         persistent := isPersistent v,
         consume := consumeExpr ctx.varMap v
     }
@@ -207,7 +212,7 @@ private def processVDecl (ctx : Context) (z : VarId) (t : IRType) (v : Expr) (b
 
 def updateVarInfoWithParams (ctx : Context) (ps : Array Param) : Context :=
   let m := ps.foldl (init := ctx.varMap) fun m p =>
-    m.insert p.x { ref := p.ty.isObj, persistent := false, consume := !p.borrow }
+    m.insert p.x { type := p.ty, persistent := false, consume := !p.borrow }
   { ctx with varMap := m }
 
 partial def visitFnBody : FnBody → Context → (FnBody × LiveVarSet)
@@ -255,7 +260,7 @@ partial def visitFnBody : FnBody → Context → (FnBody × LiveVarSet)
     match x with
     | Arg.var x =>
       let info := getVarInfo ctx x
-      if info.ref && !info.consume then (addInc ctx x b, mkLiveVarSet x) else (b, mkLiveVarSet x)
+      if info.type.isPossibleRef && !info.consume then (addInc ctx x b, mkLiveVarSet x) else (b, mkLiveVarSet x)
     | _         => (b, {})
   | b@(FnBody.jmp j xs), ctx =>
     let jLiveVars := getJPLiveVars ctx j

tests/lean/sint_basic.lean.expected.out

@@ -77,7 +77,6 @@ true
       ret x_1
     def myId8._boxed (x_1 : tagged) : tagged :=
       let x_2 : u8 := unbox x_1;
-      dec x_1;
       let x_3 : u8 := myId8 x_2;
       let x_4 : tobj := box x_3;
       ret x_4
@@ -160,7 +159,6 @@ true
       ret x_1
     def myId16._boxed (x_1 : tagged) : tagged :=
       let x_2 : u16 := unbox x_1;
-      dec x_1;
       let x_3 : u16 := myId16 x_2;
       let x_4 : tobj := box x_3;
       ret x_4