From 1bb920322d232b91f63fef5941583d85cd273eee Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Fri, 5 Apr 2019 16:51:29 -0700
Subject: [PATCH] feat(library/init/lean/compiler/constfolding): constant
 folding for `strictAnd` and `strictOr`

---
 library/init/lean/compiler/constfolding.lean | 40 +++++++++++++++-----
 tests/compiler/strictAndOr.lean              |  9 +++++
 tests/compiler/strictAndOr.lean.expected.out |  8 ++++
 3 files changed, 48 insertions(+), 9 deletions(-)
 create mode 100644 tests/compiler/strictAndOr.lean
 create mode 100644 tests/compiler/strictAndOr.lean.expected.out

diff --git a/library/init/lean/compiler/constfolding.lean b/library/init/lean/compiler/constfolding.lean
index a2e8f8d118..5ab2a92876 100644
--- a/library/init/lean/compiler/constfolding.lean
+++ b/library/init/lean/compiler/constfolding.lean
@@ -117,8 +117,36 @@ def natFoldFns : List (Name × BinFoldFn) :=
  (`Nat.decLt, foldNatDecLt),
  (`Nat.decLe, foldNatDecLe)]
 
+def getBoolLit : Expr → Option Bool
+| (Expr.const `Bool.true _)  := some true
+| (Expr.const `Bool.false _) := some false
+| _                          := none
+
+def foldStrictAnd (_ : Bool) (a₁ a₂ : Expr) : Option Expr :=
+let v₁ := getBoolLit a₁ in
+let v₂ := getBoolLit a₂ in
+match v₁, v₂ with
+| some true,  _ := a₂
+| some false, _ := a₁
+| _, some true  := a₁
+| _, some false := a₂
+| _, _          := none
+
+def foldStrictOr (_ : Bool) (a₁ a₂ : Expr) : Option Expr :=
+let v₁ := getBoolLit a₁ in
+let v₂ := getBoolLit a₂ in
+match v₁, v₂ with
+| some true,  _ := a₁
+| some false, _ := a₂
+| _, some true  := a₂
+| _, some false := a₁
+| _, _          := none
+
+def boolFoldFns : List (Name × BinFoldFn) :=
+[(`strictOr, foldStrictOr), (`strictAnd, foldStrictAnd)]
+
 def binFoldFns : List (Name × BinFoldFn) :=
-uintBinFoldFns ++ natFoldFns
+boolFoldFns ++ uintBinFoldFns ++ natFoldFns
 
 def foldNatSucc (_ : Bool) (a : Expr) : Option Expr :=
 do n   ← getNumLit a,
@@ -135,17 +163,11 @@ def unFoldFns : List (Name × UnFoldFn) :=
 [(`Nat.succ, foldNatSucc),
  (`Char.ofNat, foldCharOfNat)]
 
--- TODO(Leo): move
-private def {u} alistFind {α : Type u} (n : Name) : List (Name × α) → Option α
-| []          := none
-| ((k, v)::r) :=
-  if n = k then some v else alistFind r
-
 def findBinFoldFn (fn : Name) : Option BinFoldFn :=
-alistFind fn binFoldFns
+binFoldFns.lookup fn
 
 def findUnFoldFn (fn : Name) : Option UnFoldFn :=
-alistFind fn unFoldFns
+unFoldFns.lookup fn
 
 @[export lean.fold_bin_op_core]
 def foldBinOp (beforeErasure : Bool) (f : Expr) (a : Expr) (b : Expr) : Option Expr :=
diff --git a/tests/compiler/strictAndOr.lean b/tests/compiler/strictAndOr.lean
new file mode 100644
index 0000000000..441523d949
--- /dev/null
+++ b/tests/compiler/strictAndOr.lean
@@ -0,0 +1,9 @@
+def main : IO Unit :=
+IO.println (strictOr false false) *>
+IO.println (strictOr false true) *>
+IO.println (strictOr true false) *>
+IO.println (strictOr true true) *>
+IO.println (strictAnd false false) *>
+IO.println (strictAnd false true) *>
+IO.println (strictAnd true false) *>
+IO.println (strictAnd true true)
diff --git a/tests/compiler/strictAndOr.lean.expected.out b/tests/compiler/strictAndOr.lean.expected.out
new file mode 100644
index 0000000000..377fb6317e
--- /dev/null
+++ b/tests/compiler/strictAndOr.lean.expected.out
@@ -0,0 +1,8 @@
+false
+true
+true
+true
+false
+false
+false
+true