diff --git a/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean b/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean index a88d046e30..9654b91f47 100644 --- a/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean +++ b/src/Lean/Compiler/LCNF/Simp/ConstantFold.lean @@ -72,7 +72,7 @@ instance : Literal Nat where getLit := getNatLit mkLit := mkNatLit -partial def getStringLit (fvarId : FVarId) : CompilerM (Option String) := do +def getStringLit (fvarId : FVarId) : CompilerM (Option String) := do let some (.value (.strVal s)) ← findLetValue? fvarId | return none return s @@ -83,6 +83,18 @@ instance : Literal String where getLit := getStringLit mkLit := mkStringLit +def getBoolLit (fvarId : FVarId) : CompilerM (Option Bool) := do + let some (.const ctor [] #[]) ← findLetValue? fvarId | return none + return ctor == ``Bool.true + +def mkBoolLit (b : Bool) : FolderM LetValue := + let ctor := if b then ``Bool.true else ``Bool.false + return .const ctor [] #[] + +instance : Literal Bool where + getLit := getBoolLit + mkLit := mkBoolLit + private partial def getLitAux [Inhabited α] (fvarId : FVarId) (ofNat : Nat → α) (ofNatName : Name) : CompilerM (Option α) := do let some (.const declName _ #[.fvar fvarId]) ← findLetValue? fvarId | return none unless declName == ofNatName do return none @@ -188,6 +200,15 @@ def Folder.mkBinary [Literal α] [Literal β] [Literal γ] (folder : α → β let some arg₂ ← getLit fvarId₂ | return none mkLit <| folder arg₁ arg₂ +def Folder.mkBinaryDecisionProcedure [Literal α] [Literal β] {r : α → β → Prop} (folder : (a : α) → (b : β) → Decidable (r a b)) : Folder := fun args => do + if (← getPhase) < .mono then + return none + let #[.fvar fvarId₁, .fvar fvarId₂] := args | return none + let some arg₁ ← getLit fvarId₁ | return none + let some arg₂ ← getLit fvarId₂ | return none + let boolLit := folder arg₁ arg₂ |>.decide + mkLit boolLit + /-- Provide a folder for an operation with a left neutral element. -/ @@ -282,6 +303,26 @@ def arithmeticFolders : List (Name × Folder) := [ (``UInt64.div, Folder.first #[Folder.mkBinary UInt64.div, Folder.rightNeutral (1 : UInt64)]) ] +def relationFolders : List (Name × Folder) := [ + (``Nat.decEq, Folder.mkBinaryDecisionProcedure Nat.decEq), + (``Nat.decLt, Folder.mkBinaryDecisionProcedure Nat.decLt), + (``Nat.decLe, Folder.mkBinaryDecisionProcedure Nat.decLe), + (``UInt8.decEq, Folder.mkBinaryDecisionProcedure UInt8.decEq), + (``UInt8.decLt, Folder.mkBinaryDecisionProcedure UInt8.decLt), + (``UInt8.decLe, Folder.mkBinaryDecisionProcedure UInt8.decLe), + (``UInt16.decEq, Folder.mkBinaryDecisionProcedure UInt16.decEq), + (``UInt16.decLt, Folder.mkBinaryDecisionProcedure UInt16.decLt), + (``UInt16.decLe, Folder.mkBinaryDecisionProcedure UInt16.decLe), + (``UInt32.decEq, Folder.mkBinaryDecisionProcedure UInt32.decEq), + (``UInt32.decLt, Folder.mkBinaryDecisionProcedure UInt32.decLt), + (``UInt32.decLe, Folder.mkBinaryDecisionProcedure UInt32.decLe), + (``UInt64.decEq, Folder.mkBinaryDecisionProcedure UInt64.decEq), + (``UInt64.decLt, Folder.mkBinaryDecisionProcedure UInt64.decLt), + (``UInt64.decLe, Folder.mkBinaryDecisionProcedure UInt64.decLe), + (``Bool.decEq, Folder.mkBinaryDecisionProcedure Bool.decEq), + (``Bool.decEq, Folder.mkBinaryDecisionProcedure String.decEq) +] + /-- All string folders. -/ @@ -314,7 +355,7 @@ private def getFolder (declName : Name) : CoreM Folder := do ofExcept <| getFolderCore (← getEnv) (← getOptions) declName def builtinFolders : SMap Name Folder := - (arithmeticFolders ++ higherOrderLiteralFolders ++ stringFolders).foldl (init := {}) fun s (declName, folder) => + (arithmeticFolders ++ relationFolders ++ higherOrderLiteralFolders ++ stringFolders).foldl (init := {}) fun s (declName, folder) => s.insert declName folder structure FolderOleanEntry where