From dc71fafac1925f2a1265924a04ebc5963d014211 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Tue, 14 May 2019 18:25:54 -0700
Subject: [PATCH] feat(library/init/data/array): add `Array.binSearch`

---
 library/init/data/array/binsearch.lean   |  28 +++
 library/init/data/array/default.lean     |   1 +
 library/init/data/array/qsort.lean       |   1 +
 src/stage0/CMakeLists.txt                |   2 +-
 src/stage0/init/data/array/binsearch.cpp | 240 +++++++++++++++++++++++
 src/stage0/init/data/array/default.cpp   |   5 +-
 tests/playground/binsearch.lean          |  12 ++
 7 files changed, 287 insertions(+), 2 deletions(-)
 create mode 100644 library/init/data/array/binsearch.lean
 create mode 100644 src/stage0/init/data/array/binsearch.cpp
 create mode 100644 tests/playground/binsearch.lean

diff --git a/library/init/data/array/binsearch.lean b/library/init/data/array/binsearch.lean
new file mode 100644
index 0000000000..bc41f721f8
--- /dev/null
+++ b/library/init/data/array/binsearch.lean
@@ -0,0 +1,28 @@
+/-
+Copyright (c) 2019 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import init.data.array.basic
+
+namespace Array
+-- TODO: remove the [Inhabited α] parameters as soon as we have the tactic framework for automating proof generation and using Array.fget
+-- TODO: remove `partial` using well-founded recursion
+
+@[specialize] partial def binSearchAux {α : Type} [Inhabited α] (lt : α → α → Bool) (as : Array α) (k : α) : Nat → Nat → Option α
+| lo hi :=
+  if lo <= hi then
+    let m := (lo + hi)/2 in
+    let a := as.get m in
+    if lt a k then binSearchAux (m+1) hi
+    else if lt k a then
+      if m == 0 then none
+      else binSearchAux lo (m-1)
+    else some a
+  else none
+
+@[inline] def binSearch {α : Type} [Inhabited α] (as : Array α) (k : α) (lt : α → α → Bool) (lo := 0) (hi := as.size - 1) : Option α :=
+binSearchAux lt as k lo hi
+
+end Array
diff --git a/library/init/data/array/default.lean b/library/init/data/array/default.lean
index 043825d613..0097930324 100644
--- a/library/init/data/array/default.lean
+++ b/library/init/data/array/default.lean
@@ -6,3 +6,4 @@ Authors: Gabriel Ebner
 prelude
 import init.data.array.basic
 import init.data.array.qsort
+import init.data.array.binsearch
diff --git a/library/init/data/array/qsort.lean b/library/init/data/array/qsort.lean
index c2d4767906..3d68ce1a89 100644
--- a/library/init/data/array/qsort.lean
+++ b/library/init/data/array/qsort.lean
@@ -8,6 +8,7 @@ import init.data.array.basic
 
 namespace Array
 -- TODO: remove the [Inhabited α] parameters as soon as we have the tactic framework for automating proof generation and using Array.fget
+-- TODO: remove `partial` using well-founded recursion
 
 @[specialize] private partial def partitionAux {α : Type} [Inhabited α] (lt : α → α → Bool) (hi : Nat) (pivot : α) : Array α → Nat → Nat → Nat × Array α
 | as i j :=
diff --git a/src/stage0/CMakeLists.txt b/src/stage0/CMakeLists.txt
index dc7d00dad1..0d87826f53 100644
--- a/src/stage0/CMakeLists.txt
+++ b/src/stage0/CMakeLists.txt
@@ -1 +1 @@
-add_library (stage0 OBJECT ./init/coe.cpp ./init/control/alternative.cpp ./init/control/applicative.cpp ./init/control/combinators.cpp ./init/control/conditional.cpp ./init/control/default.cpp ./init/control/estate.cpp ./init/control/except.cpp ./init/control/functor.cpp ./init/control/id.cpp ./init/control/lift.cpp ./init/control/monad.cpp ./init/control/monadfail.cpp ./init/control/option.cpp ./init/control/reader.cpp ./init/control/state.cpp ./init/core.cpp ./init/data/array/basic.cpp ./init/data/array/default.cpp ./init/data/array/qsort.cpp ./init/data/assoclist.cpp ./init/data/basic.cpp ./init/data/bytearray/basic.cpp ./init/data/bytearray/default.cpp ./init/data/char/basic.cpp ./init/data/char/default.cpp ./init/data/default.cpp ./init/data/dlist.cpp ./init/data/fin/basic.cpp ./init/data/fin/default.cpp ./init/data/hashable.cpp ./init/data/hashmap/basic.cpp ./init/data/hashmap/default.cpp ./init/data/int/basic.cpp ./init/data/int/default.cpp ./init/data/list/basic.cpp ./init/data/list/default.cpp ./init/data/list/instances.cpp ./init/data/nat/basic.cpp ./init/data/nat/bitwise.cpp ./init/data/nat/default.cpp ./init/data/nat/div.cpp ./init/data/option/basic.cpp ./init/data/option/instances.cpp ./init/data/ordering/basic.cpp ./init/data/ordering/default.cpp ./init/data/random.cpp ./init/data/rbmap/basic.cpp ./init/data/rbmap/default.cpp ./init/data/rbtree/basic.cpp ./init/data/rbtree/default.cpp ./init/data/repr.cpp ./init/data/string/basic.cpp ./init/data/string/default.cpp ./init/data/tostring.cpp ./init/data/uint.cpp ./init/default.cpp ./init/env_ext.cpp ./init/fix.cpp ./init/io.cpp ./init/lean/compiler/constfolding.cpp ./init/lean/compiler/default.cpp ./init/lean/compiler/ir/basic.cpp ./init/lean/compiler/ir/boxing.cpp ./init/lean/compiler/ir/checker.cpp ./init/lean/compiler/ir/default.cpp ./init/lean/compiler/ir/elimdead.cpp ./init/lean/compiler/ir/format.cpp ./init/lean/compiler/ir/freevars.cpp ./init/lean/compiler/ir/livevars.cpp ./init/lean/compiler/ir/normids.cpp ./init/lean/compiler/ir/pushproj.cpp ./init/lean/compiler/ir/resetreuse.cpp ./init/lean/compiler/ir/simpcase.cpp ./init/lean/compiler/util.cpp ./init/lean/config.cpp ./init/lean/declaration.cpp ./init/lean/default.cpp ./init/lean/disjoint_set.cpp ./init/lean/elaborator.cpp ./init/lean/environment.cpp ./init/lean/expander.cpp ./init/lean/expr.cpp ./init/lean/extern.cpp ./init/lean/format.cpp ./init/lean/frontend.cpp ./init/lean/kvmap.cpp ./init/lean/level.cpp ./init/lean/message.cpp ./init/lean/modifiers.cpp ./init/lean/name.cpp ./init/lean/name_mangling.cpp ./init/lean/options.cpp ./init/lean/parser/basic.cpp ./init/lean/parser/combinators.cpp ./init/lean/parser/command.cpp ./init/lean/parser/declaration.cpp ./init/lean/parser/identifier.cpp ./init/lean/parser/level.cpp ./init/lean/parser/module.cpp ./init/lean/parser/notation.cpp ./init/lean/parser/parsec.cpp ./init/lean/parser/pratt.cpp ./init/lean/parser/rec.cpp ./init/lean/parser/stringliteral.cpp ./init/lean/parser/syntax.cpp ./init/lean/parser/term.cpp ./init/lean/parser/token.cpp ./init/lean/parser/trie.cpp ./init/lean/position.cpp ./init/lean/smap.cpp ./init/lean/trace.cpp ./init/lean/util.cpp ./init/platform.cpp ./init/util.cpp ./init/wf.cpp)
+add_library (stage0 OBJECT ./init/coe.cpp ./init/control/alternative.cpp ./init/control/applicative.cpp ./init/control/combinators.cpp ./init/control/conditional.cpp ./init/control/default.cpp ./init/control/estate.cpp ./init/control/except.cpp ./init/control/functor.cpp ./init/control/id.cpp ./init/control/lift.cpp ./init/control/monad.cpp ./init/control/monadfail.cpp ./init/control/option.cpp ./init/control/reader.cpp ./init/control/state.cpp ./init/core.cpp ./init/data/array/basic.cpp ./init/data/array/binsearch.cpp ./init/data/array/default.cpp ./init/data/array/qsort.cpp ./init/data/assoclist.cpp ./init/data/basic.cpp ./init/data/bytearray/basic.cpp ./init/data/bytearray/default.cpp ./init/data/char/basic.cpp ./init/data/char/default.cpp ./init/data/default.cpp ./init/data/dlist.cpp ./init/data/fin/basic.cpp ./init/data/fin/default.cpp ./init/data/hashable.cpp ./init/data/hashmap/basic.cpp ./init/data/hashmap/default.cpp ./init/data/int/basic.cpp ./init/data/int/default.cpp ./init/data/list/basic.cpp ./init/data/list/default.cpp ./init/data/list/instances.cpp ./init/data/nat/basic.cpp ./init/data/nat/bitwise.cpp ./init/data/nat/default.cpp ./init/data/nat/div.cpp ./init/data/option/basic.cpp ./init/data/option/instances.cpp ./init/data/ordering/basic.cpp ./init/data/ordering/default.cpp ./init/data/random.cpp ./init/data/rbmap/basic.cpp ./init/data/rbmap/default.cpp ./init/data/rbtree/basic.cpp ./init/data/rbtree/default.cpp ./init/data/repr.cpp ./init/data/string/basic.cpp ./init/data/string/default.cpp ./init/data/tostring.cpp ./init/data/uint.cpp ./init/default.cpp ./init/env_ext.cpp ./init/fix.cpp ./init/io.cpp ./init/lean/compiler/constfolding.cpp ./init/lean/compiler/default.cpp ./init/lean/compiler/ir/basic.cpp ./init/lean/compiler/ir/boxing.cpp ./init/lean/compiler/ir/checker.cpp ./init/lean/compiler/ir/default.cpp ./init/lean/compiler/ir/elimdead.cpp ./init/lean/compiler/ir/format.cpp ./init/lean/compiler/ir/freevars.cpp ./init/lean/compiler/ir/livevars.cpp ./init/lean/compiler/ir/normids.cpp ./init/lean/compiler/ir/pushproj.cpp ./init/lean/compiler/ir/resetreuse.cpp ./init/lean/compiler/ir/simpcase.cpp ./init/lean/compiler/util.cpp ./init/lean/config.cpp ./init/lean/declaration.cpp ./init/lean/default.cpp ./init/lean/disjoint_set.cpp ./init/lean/elaborator.cpp ./init/lean/environment.cpp ./init/lean/expander.cpp ./init/lean/expr.cpp ./init/lean/extern.cpp ./init/lean/format.cpp ./init/lean/frontend.cpp ./init/lean/kvmap.cpp ./init/lean/level.cpp ./init/lean/message.cpp ./init/lean/modifiers.cpp ./init/lean/name.cpp ./init/lean/name_mangling.cpp ./init/lean/options.cpp ./init/lean/parser/basic.cpp ./init/lean/parser/combinators.cpp ./init/lean/parser/command.cpp ./init/lean/parser/declaration.cpp ./init/lean/parser/identifier.cpp ./init/lean/parser/level.cpp ./init/lean/parser/module.cpp ./init/lean/parser/notation.cpp ./init/lean/parser/parsec.cpp ./init/lean/parser/pratt.cpp ./init/lean/parser/rec.cpp ./init/lean/parser/stringliteral.cpp ./init/lean/parser/syntax.cpp ./init/lean/parser/term.cpp ./init/lean/parser/token.cpp ./init/lean/parser/trie.cpp ./init/lean/position.cpp ./init/lean/smap.cpp ./init/lean/trace.cpp ./init/lean/util.cpp ./init/platform.cpp ./init/util.cpp ./init/wf.cpp)
diff --git a/src/stage0/init/data/array/binsearch.cpp b/src/stage0/init/data/array/binsearch.cpp
new file mode 100644
index 0000000000..41152b2b20
--- /dev/null
+++ b/src/stage0/init/data/array/binsearch.cpp
@@ -0,0 +1,240 @@
+// Lean compiler output
+// Module: init.data.array.binsearch
+// Imports: init.data.array.basic
+#include "runtime/object.h"
+#include "runtime/apply.h"
+typedef lean::object obj;    typedef lean::usize  usize;
+typedef lean::uint8  uint8;  typedef lean::uint16 uint16;
+typedef lean::uint32 uint32; typedef lean::uint64 uint64;
+#if defined(__clang__)
+#pragma clang diagnostic ignored "-Wunused-parameter"
+#pragma clang diagnostic ignored "-Wunused-label"
+#elif defined(__GNUC__) && !defined(__CLANG__)
+#pragma GCC diagnostic ignored "-Wunused-parameter"
+#pragma GCC diagnostic ignored "-Wunused-label"
+#pragma GCC diagnostic ignored "-Wunused-but-set-variable"
+#endif
+obj* l_Array_binSearch___rarg(obj*, obj*, obj*, obj*, obj*, obj*);
+namespace lean {
+obj* nat_sub(obj*, obj*);
+}
+obj* l_Array_binSearchAux___boxed(obj*);
+obj* l_Array_binSearchAux___rarg___boxed(obj*, obj*, obj*, obj*, obj*, obj*);
+obj* l_Array_binSearch___rarg___boxed(obj*, obj*, obj*, obj*, obj*, obj*);
+obj* l_Array_binSearch___boxed(obj*);
+obj* l_Array_binSearch(obj*);
+obj* l_Array_binSearchAux___main___rarg___boxed(obj*, obj*, obj*, obj*, obj*, obj*);
+namespace lean {
+obj* nat_add(obj*, obj*);
+}
+namespace lean {
+uint8 nat_dec_eq(obj*, obj*);
+}
+obj* l_Array_binSearchAux___main___rarg(obj*, obj*, obj*, obj*, obj*, obj*);
+obj* l_Array_binSearchAux(obj*);
+obj* l_Array_binSearchAux___main___boxed(obj*);
+namespace lean {
+uint8 nat_dec_le(obj*, obj*);
+}
+namespace lean {
+obj* nat_div(obj*, obj*);
+}
+obj* l_Array_binSearchAux___main(obj*);
+obj* l_Array_binSearchAux___rarg(obj*, obj*, obj*, obj*, obj*, obj*);
+obj* l_Array_binSearchAux___main___rarg(obj* x_0, obj* x_1, obj* x_2, obj* x_3, obj* x_4, obj* x_5) {
+_start:
+{
+uint8 x_6; 
+x_6 = lean::nat_dec_le(x_4, x_5);
+if (x_6 == 0)
+{
+obj* x_12; 
+lean::dec(x_5);
+lean::dec(x_4);
+lean::dec(x_1);
+lean::dec(x_3);
+lean::dec(x_0);
+x_12 = lean::box(0);
+return x_12;
+}
+else
+{
+obj* x_13; obj* x_14; obj* x_15; obj* x_18; obj* x_22; uint8 x_23; 
+x_13 = lean::nat_add(x_4, x_5);
+x_14 = lean::mk_nat_obj(2ul);
+x_15 = lean::nat_div(x_13, x_14);
+lean::dec(x_13);
+lean::inc(x_0);
+x_18 = lean::array_get(x_0, x_2, x_15);
+lean::inc(x_3);
+lean::inc(x_18);
+lean::inc(x_1);
+x_22 = lean::apply_2(x_1, x_18, x_3);
+x_23 = lean::unbox(x_22);
+if (x_23 == 0)
+{
+obj* x_28; uint8 x_29; 
+lean::dec(x_5);
+lean::inc(x_18);
+lean::inc(x_3);
+lean::inc(x_1);
+x_28 = lean::apply_2(x_1, x_3, x_18);
+x_29 = lean::unbox(x_28);
+if (x_29 == 0)
+{
+obj* x_35; 
+lean::dec(x_15);
+lean::dec(x_4);
+lean::dec(x_1);
+lean::dec(x_3);
+lean::dec(x_0);
+x_35 = lean::alloc_cnstr(1, 1, 0);
+lean::cnstr_set(x_35, 0, x_18);
+return x_35;
+}
+else
+{
+obj* x_37; uint8 x_38; 
+lean::dec(x_18);
+x_37 = lean::mk_nat_obj(0ul);
+x_38 = lean::nat_dec_eq(x_15, x_37);
+if (x_38 == 0)
+{
+obj* x_39; obj* x_40; 
+x_39 = lean::mk_nat_obj(1ul);
+x_40 = lean::nat_sub(x_15, x_39);
+lean::dec(x_15);
+x_5 = x_40;
+goto _start;
+}
+else
+{
+obj* x_48; 
+lean::dec(x_15);
+lean::dec(x_4);
+lean::dec(x_1);
+lean::dec(x_3);
+lean::dec(x_0);
+x_48 = lean::box(0);
+return x_48;
+}
+}
+}
+else
+{
+obj* x_51; obj* x_52; 
+lean::dec(x_18);
+lean::dec(x_4);
+x_51 = lean::mk_nat_obj(1ul);
+x_52 = lean::nat_add(x_15, x_51);
+lean::dec(x_15);
+x_4 = x_52;
+goto _start;
+}
+}
+}
+}
+obj* l_Array_binSearchAux___main(obj* x_0) {
+_start:
+{
+obj* x_1; 
+x_1 = lean::alloc_closure(reinterpret_cast<void*>(l_Array_binSearchAux___main___rarg___boxed), 6, 0);
+return x_1;
+}
+}
+obj* l_Array_binSearchAux___main___rarg___boxed(obj* x_0, obj* x_1, obj* x_2, obj* x_3, obj* x_4, obj* x_5) {
+_start:
+{
+obj* x_6; 
+x_6 = l_Array_binSearchAux___main___rarg(x_0, x_1, x_2, x_3, x_4, x_5);
+lean::dec(x_2);
+return x_6;
+}
+}
+obj* l_Array_binSearchAux___main___boxed(obj* x_0) {
+_start:
+{
+obj* x_1; 
+x_1 = l_Array_binSearchAux___main(x_0);
+lean::dec(x_0);
+return x_1;
+}
+}
+obj* l_Array_binSearchAux___rarg(obj* x_0, obj* x_1, obj* x_2, obj* x_3, obj* x_4, obj* x_5) {
+_start:
+{
+obj* x_6; 
+x_6 = l_Array_binSearchAux___main___rarg(x_0, x_1, x_2, x_3, x_4, x_5);
+return x_6;
+}
+}
+obj* l_Array_binSearchAux(obj* x_0) {
+_start:
+{
+obj* x_1; 
+x_1 = lean::alloc_closure(reinterpret_cast<void*>(l_Array_binSearchAux___rarg___boxed), 6, 0);
+return x_1;
+}
+}
+obj* l_Array_binSearchAux___rarg___boxed(obj* x_0, obj* x_1, obj* x_2, obj* x_3, obj* x_4, obj* x_5) {
+_start:
+{
+obj* x_6; 
+x_6 = l_Array_binSearchAux___rarg(x_0, x_1, x_2, x_3, x_4, x_5);
+lean::dec(x_2);
+return x_6;
+}
+}
+obj* l_Array_binSearchAux___boxed(obj* x_0) {
+_start:
+{
+obj* x_1; 
+x_1 = l_Array_binSearchAux(x_0);
+lean::dec(x_0);
+return x_1;
+}
+}
+obj* l_Array_binSearch___rarg(obj* x_0, obj* x_1, obj* x_2, obj* x_3, obj* x_4, obj* x_5) {
+_start:
+{
+obj* x_6; 
+x_6 = l_Array_binSearchAux___main___rarg(x_0, x_3, x_1, x_2, x_4, x_5);
+return x_6;
+}
+}
+obj* l_Array_binSearch(obj* x_0) {
+_start:
+{
+obj* x_1; 
+x_1 = lean::alloc_closure(reinterpret_cast<void*>(l_Array_binSearch___rarg___boxed), 6, 0);
+return x_1;
+}
+}
+obj* l_Array_binSearch___rarg___boxed(obj* x_0, obj* x_1, obj* x_2, obj* x_3, obj* x_4, obj* x_5) {
+_start:
+{
+obj* x_6; 
+x_6 = l_Array_binSearch___rarg(x_0, x_1, x_2, x_3, x_4, x_5);
+lean::dec(x_1);
+return x_6;
+}
+}
+obj* l_Array_binSearch___boxed(obj* x_0) {
+_start:
+{
+obj* x_1; 
+x_1 = l_Array_binSearch(x_0);
+lean::dec(x_0);
+return x_1;
+}
+}
+obj* initialize_init_data_array_basic(obj*);
+static bool _G_initialized = false;
+obj* initialize_init_data_array_binsearch(obj* w) {
+ if (_G_initialized) return w;
+ _G_initialized = true;
+if (io_result_is_error(w)) return w;
+w = initialize_init_data_array_basic(w);
+if (io_result_is_error(w)) return w;
+return w;
+}
diff --git a/src/stage0/init/data/array/default.cpp b/src/stage0/init/data/array/default.cpp
index 18ce3026d5..81ae008b3d 100644
--- a/src/stage0/init/data/array/default.cpp
+++ b/src/stage0/init/data/array/default.cpp
@@ -1,6 +1,6 @@
 // Lean compiler output
 // Module: init.data.array.default
-// Imports: init.data.array.basic init.data.array.qsort
+// Imports: init.data.array.basic init.data.array.qsort init.data.array.binsearch
 #include "runtime/object.h"
 #include "runtime/apply.h"
 typedef lean::object obj;    typedef lean::usize  usize;
@@ -16,6 +16,7 @@ typedef lean::uint32 uint32; typedef lean::uint64 uint64;
 #endif
 obj* initialize_init_data_array_basic(obj*);
 obj* initialize_init_data_array_qsort(obj*);
+obj* initialize_init_data_array_binsearch(obj*);
 static bool _G_initialized = false;
 obj* initialize_init_data_array_default(obj* w) {
  if (_G_initialized) return w;
@@ -25,5 +26,7 @@ w = initialize_init_data_array_basic(w);
 if (io_result_is_error(w)) return w;
 w = initialize_init_data_array_qsort(w);
 if (io_result_is_error(w)) return w;
+w = initialize_init_data_array_binsearch(w);
+if (io_result_is_error(w)) return w;
 return w;
 }
diff --git a/tests/playground/binsearch.lean b/tests/playground/binsearch.lean
new file mode 100644
index 0000000000..039084676a
--- /dev/null
+++ b/tests/playground/binsearch.lean
@@ -0,0 +1,12 @@
+def mkAssocArray : Nat → Array (Nat × Bool) → Array (Nat × Bool)
+| 0     as := as
+| (i+1) as := mkAssocArray i (as.push (i, i % 2 == 0))
+
+def main (xs : List String) : IO Unit :=
+do
+let n  := xs.head.toNat,
+let as := mkAssocArray n Array.empty,
+let as := as.qsort (λ a b, a.1 < b.1),
+(2*n).mfor $ λ i, do
+  let entry := as.binSearch (i, false) (λ a b, a.1 < b.1),
+  IO.println (">> " ++ toString i ++ " ==> " ++ toString entry)