From fc8f29034778cd56c6979f4979475e49fb02ac97 Mon Sep 17 00:00:00 2001 From: Cameron Zwarich Date: Mon, 19 May 2025 17:38:38 -0700 Subject: [PATCH] feat: support native literals of size unsigned integer types (#8409) This PR adds support to LCNF for native UInt8/UInt16/UInt32/UInt64 literals. --- src/Lean/Compiler/IR/EmitC.lean | 5 ++++- src/Lean/Compiler/IR/ToIR.lean | 4 ++++ src/Lean/Compiler/LCNF/Basic.lean | 10 +++++++++- src/Lean/Compiler/LCNF/ElimDeadBranches.lean | 4 ++++ src/Lean/Compiler/LCNF/InferType.lean | 4 ++++ src/Lean/Compiler/LCNF/PrettyPrinter.lean | 2 +- src/Lean/Compiler/LCNF/Simp/ConstantFold.lean | 17 ++++++++++++++++- 7 files changed, 42 insertions(+), 4 deletions(-) diff --git a/src/Lean/Compiler/IR/EmitC.lean b/src/Lean/Compiler/IR/EmitC.lean index 52110585ce..bf6489bdd4 100644 --- a/src/Lean/Compiler/IR/EmitC.lean +++ b/src/Lean/Compiler/IR/EmitC.lean @@ -495,7 +495,10 @@ def emitNumLit (t : IRType) (v : Nat) : M Unit := do else emit "lean_cstr_to_nat(\""; emit v; emit "\")" else - emit v + if v < UInt32.size then + emit v + else + emit v; emit "ULL" def emitLit (z : VarId) (t : IRType) (v : LitVal) : M Unit := do emitLhs z; diff --git a/src/Lean/Compiler/IR/ToIR.lean b/src/Lean/Compiler/IR/ToIR.lean index ac7a71188b..6ca8bd730a 100644 --- a/src/Lean/Compiler/IR/ToIR.lean +++ b/src/Lean/Compiler/IR/ToIR.lean @@ -67,6 +67,10 @@ def lowerLitValue (v : LCNF.LitValue) : LitVal := match v with | .nat n => .num n | .str s => .str s + | .uint8 v => .num (UInt8.toNat v) + | .uint16 v => .num (UInt16.toNat v) + | .uint32 v => .num (UInt32.toNat v) + | .uint64 v => .num (UInt64.toNat v) -- TODO: This should be cached. def lowerEnumToScalarType (name : Name) : M (Option IRType) := do diff --git a/src/Lean/Compiler/LCNF/Basic.lean b/src/Lean/Compiler/LCNF/Basic.lean index d1eb4359f5..2da4e53f03 100644 --- a/src/Lean/Compiler/LCNF/Basic.lean +++ b/src/Lean/Compiler/LCNF/Basic.lean @@ -36,12 +36,20 @@ def Param.toExpr (p : Param) : Expr := inductive LitValue where | nat (val : Nat) | str (val : String) - -- TODO: add constructors for `Int`, `Float`, `UInt` ... + | uint8 (val : UInt8) + | uint16 (val : UInt16) + | uint32 (val : UInt32) + | uint64 (val : UInt64) + -- TODO: add constructors for `Int`, `Float`, `USize` ... deriving Inhabited, BEq, Hashable def LitValue.toExpr : LitValue → Expr | .nat v => .lit (.natVal v) | .str v => .lit (.strVal v) + | .uint8 v => .app (.const ``UInt8.ofNat []) (.lit (.natVal (UInt8.toNat v))) + | .uint16 v => .app (.const ``UInt16.ofNat []) (.lit (.natVal (UInt16.toNat v))) + | .uint32 v => .app (.const ``UInt32.ofNat []) (.lit (.natVal (UInt32.toNat v))) + | .uint64 v => .app (.const ``UInt64.ofNat []) (.lit (.natVal (UInt64.toNat v))) inductive Arg where | erased diff --git a/src/Lean/Compiler/LCNF/ElimDeadBranches.lean b/src/Lean/Compiler/LCNF/ElimDeadBranches.lean index ae3dab8945..227194df39 100644 --- a/src/Lean/Compiler/LCNF/ElimDeadBranches.lean +++ b/src/Lean/Compiler/LCNF/ElimDeadBranches.lean @@ -174,6 +174,10 @@ def ofLCNFLit : LCNF.LitValue → Value | .nat n => ofNat n -- TODO: We could make this much more precise but the payoff is questionable | .str .. => .top +| .uint8 v => ofNat (UInt8.toNat v) +| .uint16 v => ofNat (UInt16.toNat v) +| .uint32 v => ofNat (UInt32.toNat v) +| .uint64 v => ofNat (UInt64.toNat v) partial def proj : Value → Nat → Value | .ctor _ vs , i => vs.getD i bot diff --git a/src/Lean/Compiler/LCNF/InferType.lean b/src/Lean/Compiler/LCNF/InferType.lean index 546220459b..f6b0aad9dd 100644 --- a/src/Lean/Compiler/LCNF/InferType.lean +++ b/src/Lean/Compiler/LCNF/InferType.lean @@ -105,6 +105,10 @@ def inferLitValueType (value : LitValue) : Expr := match value with | .nat .. => mkConst ``Nat | .str .. => mkConst ``String + | .uint8 .. => mkConst ``UInt8 + | .uint16 .. => mkConst ``UInt16 + | .uint32 .. => mkConst ``UInt32 + | .uint64 .. => mkConst ``UInt64 mutual partial def inferArgType (arg : Arg) : InferTypeM Expr := diff --git a/src/Lean/Compiler/LCNF/PrettyPrinter.lean b/src/Lean/Compiler/LCNF/PrettyPrinter.lean index dc9f1487c5..2032787f44 100644 --- a/src/Lean/Compiler/LCNF/PrettyPrinter.lean +++ b/src/Lean/Compiler/LCNF/PrettyPrinter.lean @@ -58,7 +58,7 @@ def ppArgs (args : Array Arg) : M Format := do def ppLitValue (lit : LitValue) : M Format := do match lit with - | .nat v => return format v + | .nat v | .uint8 v | .uint16 v | .uint32 v | .uint64 v => return format v | .str v => return format (repr v) def ppLetValue (e : LetValue) : M Format := do diff --git a/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean b/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean index 45369b76f6..0dcf52b684 100644 --- a/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean +++ b/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean @@ -308,6 +308,14 @@ def higherOrderLiteralFolders : List (Name × Folder) := [ def Folder.mulShift [Literal α] [BEq α] (shiftLeft : Name) (pow2 : α → α) (log2 : α → α) : Folder := Folder.first #[Folder.mulLhsShift shiftLeft pow2 log2, Folder.mulRhsShift shiftLeft pow2 log2] +/-- +Folder for ofNat operations on fixed-sized integer types. +-/ +def Folder.ofNat (f : Nat → LitValue) (args : Array Arg): FolderM (Option LetValue) := do + let #[.fvar fvarId] := args | return none + let some value ← getNatLit fvarId | return none + return some (.lit (f value)) + /-- All arithmetic folders. -/ @@ -355,6 +363,13 @@ def relationFolders : List (Name × Folder) := [ (``Bool.decEq, Folder.mkBinaryDecisionProcedure String.decEq) ] +def conversionFolders : List (Name × Folder) := [ + (``UInt8.ofNat, Folder.ofNat (fun v => .uint8 (UInt8.ofNat v))), + (``UInt16.ofNat, Folder.ofNat (fun v => .uint16 (UInt16.ofNat v))), + (``UInt32.ofNat, Folder.ofNat (fun v => .uint32 (UInt32.ofNat v))), + (``UInt64.ofNat, Folder.ofNat (fun v => .uint64 (UInt64.ofNat v))), +] + /-- All string folders. -/ @@ -387,7 +402,7 @@ private def getFolder (declName : Name) : CoreM Folder := do ofExcept <| getFolderCore (← getEnv) (← getOptions) declName def builtinFolders : SMap Name Folder := - (arithmeticFolders ++ relationFolders ++ higherOrderLiteralFolders ++ stringFolders).foldl (init := {}) fun s (declName, folder) => + (arithmeticFolders ++ relationFolders ++ conversionFolders ++ higherOrderLiteralFolders ++ stringFolders).foldl (init := {}) fun s (declName, folder) => s.insert declName folder structure FolderOleanEntry where