more

Author: hargoniX

src/Lean/Compiler/IR/ToIR.lean

@@ -20,17 +20,25 @@ public section
 namespace Lean.IR
 
 open Lean.Compiler (LCNF.Alt LCNF.Arg LCNF.Code LCNF.Decl LCNF.DeclValue LCNF.LCtx LCNF.LetDecl
-                    LCNF.LetValue LCNF.LitValue LCNF.Param LCNF.getMonoDecl?)
+  LCNF.LetValue LCNF.LitValue LCNF.Param LCNF.getMonoDecl? LCNF.FVarSubst LCNF.MonadFVarSubst
+  LCNF.MonadFVarSubstState LCNF.addSubst LCNF.normLetValue)
 
 namespace ToIR
 
 structure BuilderState where
   vars : Std.HashMap FVarId Arg := {}
   joinPoints : Std.HashMap FVarId JoinPointId := {}
   nextId : Nat := 1
+  subst : LCNF.FVarSubst := {}
 
 abbrev M := StateRefT BuilderState CoreM
 
+instance : LCNF.MonadFVarSubst M false where
+  getSubst := return (← get).subst
+
+instance : LCNF.MonadFVarSubstState M where
+  modifySubst f := modify fun s => { s with subst := f s.subst }
+
 def M.run (x : M α) : CoreM α := do
   x.run' {}
 
@@ -119,41 +127,57 @@ partial def lowerCode (c : LCNF.Code) : M FnBody := do
     let joinPointId ← getJoinPointValue fvarId
     return .jmp joinPointId (← args.mapM lowerArg)
   | .cases cases =>
-    -- `casesOn` for inductive predicates should have already been expanded.
-    let .var varId := (← getFVarValue cases.discr) | unreachable!
-    return .case cases.typeName
-                 varId
-                 (← nameToIRType cases.typeName)
-                 (← cases.alts.mapM (lowerAlt varId))
+    if let some info ← hasTrivialStructure? cases.typeName then
+      assert! cases.alts.size == 1
+      let .alt ctorName ps k := cases.alts[0]! | unreachable!
+      assert! ctorName == info.ctorName
+      assert! info.fieldIdx < ps.size
+      let p := ps[info.fieldIdx]!
+      LCNF.addSubst p.fvarId (.fvar cases.discr)
+      lowerCode k
+    else
+      -- `casesOn` for inductive predicates should have already been expanded.
+      let .var varId := (← getFVarValue cases.discr) | unreachable!
+      return .case cases.typeName
+                   varId
+                   (← nameToIRType cases.typeName)
+                   (← cases.alts.mapM (lowerAlt varId))
   | .return fvarId =>
     return .ret (← getFVarValue fvarId)
   | .unreach .. => return .unreachable
   | .fun .. => panic! "all local functions should be λ-lifted"
 
 partial def lowerLet (decl : LCNF.LetDecl) (k : LCNF.Code) : M FnBody := do
-  match decl.value with
+  let value ← LCNF.normLetValue decl.value
+  match value with
   | .lit litValue =>
     let var ← bindVar decl.fvarId
     let ⟨litValue, type⟩ := lowerLitValue litValue
     return .vdecl var type (.lit litValue) (← lowerCode k)
   | .proj typeName i fvarId =>
-    match (← getFVarValue fvarId) with
-    | .var varId =>
-      let some (.inductInfo { ctors := [ctorName], .. }) := (← Lean.getEnv).find? typeName
-        | panic! "projection of non-structure type"
-      let ⟨ctorInfo, fields⟩ ← getCtorLayout ctorName
-      let ⟨result, type⟩ := lowerProj varId ctorInfo fields[i]!
-      match result with
-      | .expr e =>
-        let var ← bindVar decl.fvarId
-        return .vdecl var type e (← lowerCode k)
+    if let some info ← hasTrivialStructure? typeName then
+      if info.fieldIdx == i then
+        LCNF.addSubst decl.fvarId (.fvar fvarId)
+      else
+        bindErased decl.fvarId
+      lowerCode k
+    else
+      match (← getFVarValue fvarId) with
+      | .var varId =>
+        let some (.inductInfo { ctors := [ctorName], .. }) := (← Lean.getEnv).find? typeName
+          | panic! "projection of non-structure type"
+        let ⟨ctorInfo, fields⟩ ← getCtorLayout ctorName
+        let ⟨result, type⟩ := lowerProj varId ctorInfo fields[i]!
+        match result with
+        | .expr e =>
+          let var ← bindVar decl.fvarId
+          return .vdecl var type e (← lowerCode k)
+        | .erased | .void =>
+          bindErased decl.fvarId
+          lowerCode k
       | .erased =>
         bindErased decl.fvarId
         lowerCode k
-      | .void => unreachable!
-    | .erased =>
-      bindErased decl.fvarId
-      lowerCode k
   | .const name _ args =>
     let irArgs ← args.mapM lowerArg
     if let some decl ← findDecl name then
@@ -163,43 +187,48 @@ partial def lowerLet (decl : LCNF.LetDecl) (k : LCNF.Code) : M FnBody := do
     let env ← Lean.getEnv
     match env.find? name with
     | some (.ctorInfo ctorVal) =>
-      let type ← nameToIRType ctorVal.induct
-      if type.isScalar then
-        let var ← bindVar decl.fvarId
-        return .vdecl var type (.lit (.num ctorVal.cidx)) (← lowerCode k)
-
-      let ⟨ctorInfo, fields⟩ ← getCtorLayout name
-      let irArgs := irArgs.extract (start := ctorVal.numParams)
-      if irArgs.size != fields.size then
-        -- An overapplied constructor arises from compiler
-        -- transformations on unreachable code
-        return .unreachable
-
-      let objArgs : Array Arg ← do
-        let mut result : Array Arg := #[]
-        for h : i in *...fields.size do
-          match fields[i] with
-          | .object .. =>
-            result := result.push irArgs[i]!
-          | .usize .. | .scalar .. | .erased | .void => pure ()
-        pure result
-      let objVar ← bindVar decl.fvarId
-      let rec lowerNonObjectFields (_ : Unit) : M FnBody :=
-        let rec loop (i : Nat) : M FnBody := do
-          match irArgs[i]? with
-          | some (.var varId) =>
-            match fields[i]! with
-            | .usize usizeIdx =>
-              let k ← loop (i + 1)
-              return .uset objVar usizeIdx varId k
-            | .scalar _ offset argType =>
-              let k ← loop (i + 1)
-              return .sset objVar (ctorInfo.size + ctorInfo.usize) offset varId argType k
-            | .object .. | .erased | .void => loop (i + 1)
-          | some .erased => loop (i + 1)
-          | none => lowerCode k
-        loop 0
-      return .vdecl objVar ctorInfo.type (.ctor ctorInfo objArgs) (← lowerNonObjectFields ())
+      if let some info ← hasTrivialStructure? ctorVal.induct then
+        let arg := args[info.numParams + info.fieldIdx]!
+        LCNF.addSubst decl.fvarId arg
+        lowerCode k
+      else
+        let type ← nameToIRType ctorVal.induct
+        if type.isScalar then
+          let var ← bindVar decl.fvarId
+          return .vdecl var type (.lit (.num ctorVal.cidx)) (← lowerCode k)
+
+        let ⟨ctorInfo, fields⟩ ← getCtorLayout name
+        let irArgs := irArgs.extract (start := ctorVal.numParams)
+        if irArgs.size != fields.size then
+          -- An overapplied constructor arises from compiler
+          -- transformations on unreachable code
+          return .unreachable
+
+        let objArgs : Array Arg ← do
+          let mut result : Array Arg := #[]
+          for h : i in *...fields.size do
+            match fields[i] with
+            | .object .. =>
+              result := result.push irArgs[i]!
+            | .usize .. | .scalar .. | .erased | .void => pure ()
+          pure result
+        let objVar ← bindVar decl.fvarId
+        let rec lowerNonObjectFields (_ : Unit) : M FnBody :=
+          let rec loop (i : Nat) : M FnBody := do
+            match irArgs[i]? with
+            | some (.var varId) =>
+              match fields[i]! with
+              | .usize usizeIdx =>
+                let k ← loop (i + 1)
+                return .uset objVar usizeIdx varId k
+              | .scalar _ offset argType =>
+                let k ← loop (i + 1)
+                return .sset objVar (ctorInfo.size + ctorInfo.usize) offset varId argType k
+              | .object .. | .erased | .void => loop (i + 1)
+            | some .erased => loop (i + 1)
+            | none => lowerCode k
+          loop 0
+        return .vdecl objVar ctorInfo.type (.ctor ctorInfo objArgs) (← lowerNonObjectFields ())
     | some (.defnInfo ..) | some (.opaqueInfo ..) =>
       mkFap name irArgs
     | some (.axiomInfo ..) | .some (.quotInfo ..) | .some (.inductInfo ..) | .some (.thmInfo ..) =>

src/Lean/Compiler/IR/ToIRType.lean

@@ -17,7 +17,9 @@ public section
 namespace Lean
 namespace IR
 
-open Lean.Compiler (LCNF.CacheExtension LCNF.isTypeFormerType LCNF.toLCNFType LCNF.toMonoType)
+open Lean.Compiler (LCNF.CacheExtension LCNF.isTypeFormerType LCNF.toLCNFType LCNF.toMonoType
+  LCNF.TrivialStructureInfo LCNF.getOtherDeclBaseType LCNF.getParamTypes LCNF.instantiateForall
+  LCNF.Irrelevant.hasTrivialStructure?)
 
 def irTypeForEnum (numCtors : Nat) : IRType :=
   if numCtors == 1 then
@@ -34,6 +36,31 @@ def irTypeForEnum (numCtors : Nat) : IRType :=
 builtin_initialize irTypeExt : LCNF.CacheExtension Name IRType ←
   LCNF.CacheExtension.register
 
+builtin_initialize trivialStructureInfoExt :
+    LCNF.CacheExtension Name (Option LCNF.TrivialStructureInfo) ←
+  LCNF.CacheExtension.register
+
+def getRelevantCtorFields (ctorName : Name) : CoreM (Array Bool) := do
+  let .ctorInfo info ← getConstInfo ctorName | unreachable!
+  Meta.MetaM.run' do
+    Meta.forallTelescopeReducing info.type fun xs _ => do
+      let mut result := #[]
+      for x in xs[info.numParams...*] do
+        let type ← Meta.inferType x
+        let cond := !(← Meta.isProp type <||> Meta.isTypeFormerType type <||> pure type.isRealWorld)
+        result := result.push cond
+      return result
+
+/--
+The idea of this function is the same as in `ToMono`, however the notion of "irrelevancy" has
+changed because we now have the `void` type which can only be erased in impure context and thus at
+earliest at the conversion from mono to IR.
+-/
+def hasTrivialStructure? (declName : Name) : CoreM (Option LCNF.TrivialStructureInfo) := do
+  let irrelevantType type :=
+    Meta.isProp type <||> Meta.isTypeFormerType type <||> pure type.isRealWorld
+  LCNF.Irrelevant.hasTrivialStructure? trivialStructureInfoExt irrelevantType declName
+
 def nameToIRType (name : Name) : CoreM IRType := do
   match (← irTypeExt.find? name) with
   | some type => return type
@@ -86,13 +113,13 @@ private def isAnyProducingType (type : Lean.Expr) : Bool :=
   | .forallE _ _ b _ => isAnyProducingType b
   | _ => false
 
-def toIRType (type : Lean.Expr) : CoreM IRType := do
+partial def toIRType (type : Lean.Expr) : CoreM IRType := do
   match type with
-  | .const name _ => nameToIRType name
+  | .const name _ => visitApp name #[]
   | .app .. =>
     -- All mono types are in headBeta form.
     let .const name _ := type.getAppFn | unreachable!
-    nameToIRType name
+    visitApp name type.getAppArgs
   | .forallE _ _ b _ =>
     -- Type formers are erased, but can be used polymorphically as
     -- an arrow type producing `lcAny`. The runtime representation of
@@ -104,6 +131,14 @@ def toIRType (type : Lean.Expr) : CoreM IRType := do
       return .object
   | .mdata _ b => toIRType b
   | _ => unreachable!
+where
+  visitApp (declName : Name) (args : Array Lean.Expr) : CoreM IRType := do
+    if let some info ← hasTrivialStructure? declName then
+      let ctorType ← LCNF.getOtherDeclBaseType info.ctorName []
+      let monoType ← LCNF.toMonoType (LCNF.getParamTypes (← LCNF.instantiateForall ctorType args[*...info.numParams]))[info.fieldIdx]!
+      toIRType monoType
+    else
+      nameToIRType declName
 
 inductive CtorFieldInfo where
   | erased

src/Lean/Compiler/LCNF.lean

@@ -44,5 +44,6 @@ public import Lean.Compiler.LCNF.Closure
 public import Lean.Compiler.LCNF.LambdaLifting
 public import Lean.Compiler.LCNF.ReduceArity
 public import Lean.Compiler.LCNF.Probing
+public import Lean.Compiler.LCNF.Irrelevant
 
 public section

src/Lean/Compiler/LCNF/Irrelevant.lean

@@ -0,0 +1,71 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Henrik Böving
+-/
+module
+
+prelude
+public import Lean.Compiler.LCNF.CompilerM
+import Lean.Compiler.LCNF.BaseTypes
+import Lean.Compiler.LCNF.Util
+
+namespace Lean.Compiler.LCNF
+
+/--
+Given a constructor, return a bitmask `m` s.t. `m[i]` is true if field `i` is
+computationally relevant.
+-/
+def getRelevantCtorFields (ctorName : Name) (trivialType : Expr → MetaM Bool) :
+    CoreM (Array Bool) := do
+  let .ctorInfo info ← getConstInfo ctorName | unreachable!
+  Meta.MetaM.run' do
+    Meta.forallTelescopeReducing info.type fun xs _ => do
+      let mut result := #[]
+      for x in xs[info.numParams...*] do
+        let type ← Meta.inferType x
+        result := result.push !(← trivialType type)
+      return result
+
+/--
+We say a structure has a trivial structure if it has not builtin support in the runtime,
+it has only one constructor, and this constructor has only one relevant field.
+-/
+public structure TrivialStructureInfo where
+  ctorName  : Name
+  numParams : Nat
+  fieldIdx  : Nat
+  deriving Inhabited, Repr
+
+/--
+Return `some fieldIdx` if `declName` is the name of an inductive datatype s.t.
+- It does not have builtin support in the runtime.
+- It has only one constructor.
+- This constructor has only one computationally relevant field.
+-/
+public def Irrelevant.hasTrivialStructure?
+    (cacheExt : CacheExtension Name (Option TrivialStructureInfo))
+    (trivialType : Expr → MetaM Bool) (declName : Name) : CoreM (Option TrivialStructureInfo) := do
+  match (← cacheExt.find? declName) with
+  | some info? => return info?
+  | none =>
+    let info? ← fillCache
+    cacheExt.insert declName info?
+    return info?
+where fillCache : CoreM (Option TrivialStructureInfo) := do
+  if isRuntimeBuiltinType declName then return none
+  let .inductInfo info ← getConstInfo declName | return none
+  if info.isUnsafe || info.isRec then return none
+  let [ctorName] := info.ctors | return none
+  let ctorType ← getOtherDeclBaseType ctorName []
+  if ctorType.isErased then return none
+  let mask ← getRelevantCtorFields ctorName trivialType
+  let mut result := none
+  for h : i in *...mask.size do
+    if mask[i] then
+      if result.isSome then return none
+      result := some { ctorName, fieldIdx := i, numParams := info.numParams }
+  return result
+
+
+end Lean.Compiler.LCNF