From 9382b8120bbf345b7dc40dcbc855fcea65b02984 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Wed, 18 Mar 2020 05:42:15 -0700
Subject: [PATCH] [sledge] Improve IArray.map_adjacent

Summary:
- Rename to `combine_adjacent` for clarity

- Reimplement without needing to filter out a dummy value, by keeping
  track of the numer of adjacent pairs that have been combined
  together

- Use this counting to replace the probably-confusing lazy copy with
  explicit eager copy at the right time

- Remove the now-unneeded dummy value argument

Reviewed By: jvillard

Differential Revision: D20482752

fbshipit-source-id: db80bbbe3
---
 sledge/lib/import/IArray.ml  | 35 ++++++++++++++++-------------------
 sledge/lib/import/IArray.mli |  2 +-
 sledge/lib/term.ml           |  2 +-
 3 files changed, 18 insertions(+), 21 deletions(-)

diff --git a/sledge/lib/import/IArray.ml b/sledge/lib/import/IArray.ml
index d0bc26af2..4255686c9 100644
--- a/sledge/lib/import/IArray.ml
+++ b/sledge/lib/import/IArray.ml
@@ -34,26 +34,23 @@ let to_array = a
 let pp sep pp_elt fs v = List.pp sep pp_elt fs (to_list v)
 let concat_map x ~f = v (Array.concat_map (a x) ~f:(fun y -> a (f y)))
 
-let map_adjacent ~f dummy xs_v =
-  let xs0 = a xs_v in
-  let copy_xs = lazy (Array.copy xs0) in
-  let n = Array.length xs0 - 1 in
-  let rec map_adjacent_ i xs =
-    if i < n then
-      let xs =
-        match f xs.(i) xs.(i + 1) with
-        | None -> xs
-        | Some x ->
-            let xs = Lazy.force copy_xs in
-            xs.(i) <- dummy ;
-            xs.(i + 1) <- x ;
-            xs
-      in
-      map_adjacent_ (i + 1) xs
-    else if xs == xs0 then xs
-    else Array.filter xs ~f:(fun x -> not (dummy == x))
+let combine_adjacent ~f xs_v =
+  let xs = a xs_v in
+  let n = Array.length xs - 1 in
+  let rec combine_adjacent_ j i xs =
+    if i < n then (
+      match f xs.(i - j) xs.(i + 1) with
+      | None ->
+          if j != 0 then xs.(i + 1 - j) <- xs.(i + 1) ;
+          combine_adjacent_ j (i + 1) xs
+      | Some x ->
+          let xs = if j = 0 then Array.copy xs else xs in
+          xs.(i - j) <- x ;
+          combine_adjacent_ (j + 1) (i + 1) xs )
+    else if j = 0 then xs
+    else Array.sub xs ~pos:0 ~len:(n + 1 - j)
   in
-  v (map_adjacent_ 0 xs0)
+  v (combine_adjacent_ 0 0 xs)
 
 let create ~len x = v (Array.create ~len x)
 let empty = v [||]
diff --git a/sledge/lib/import/IArray.mli b/sledge/lib/import/IArray.mli
index f9d23cc55..65a747484 100644
--- a/sledge/lib/import/IArray.mli
+++ b/sledge/lib/import/IArray.mli
@@ -137,7 +137,7 @@ val concat_map : 'a t -> f:('a -> 'b t) -> 'b t
 
 (* val concat_mapi : 'a t -> f:(int -> 'a -> 'b t) -> 'b t *)
 
-val map_adjacent : f:('a -> 'a -> 'a option) -> 'a -> 'a t -> 'a t
+val combine_adjacent : f:('a -> 'a -> 'a option) -> 'a t -> 'a t
 
 (* val partition_tf : 'a t -> f:('a -> bool) -> 'a t * 'a t *)
 (* val partitioni_tf : 'a t -> f:(int -> 'a -> bool) -> 'a t * 'a t *)
diff --git a/sledge/lib/term.ml b/sledge/lib/term.ml
index f33541d4c..1b412bd67 100644
--- a/sledge/lib/term.ml
+++ b/sledge/lib/term.ml
@@ -766,7 +766,7 @@ and simp_concat xs =
     | _ -> None
   in
   let xs = flatten xs in
-  let xs = IArray.map_adjacent empty_agg xs ~f:simp_adjacent in
+  let xs = IArray.combine_adjacent ~f:simp_adjacent xs in
   (if IArray.length xs = 1 then IArray.get xs 0 else ApN (Concat, xs))
   |>
   [%Trace.retn fun {pf} -> pf "%a" pp]