[sledge] Strengthen and simplify canonizer for Extract terms

Summary:
Formulate the canonizer for Extract from Concat terms uniformly as a
concatenation of extracts.

Reviewed By: jvillard

Differential Revision: D20303064

fbshipit-source-id: a45bc45dd

sledge/src/import/vector.ml

@@ -101,6 +101,14 @@ let fold_map x ~init ~f =
   let s, x = Array.fold_map (a x) ~init ~f in
   (s, v x)
 
+let fold_map_until xs ~init ~f ~finish =
+  With_return.with_return (fun {return} ->
+      finish
+        (fold_map xs ~init ~f:(fun s x ->
+             match (f s x : _ Continue_or_stop.t) with
+             | Continue x -> x
+             | Stop x -> return x )) )
+
 let concat xs = v (Array.concat (al xs))
 let copy x = v (Array.copy (a x))
 let sub ~pos ~len x = v (Array.sub ~pos ~len (a x))

sledge/src/import/vector.mli

@@ -108,6 +108,14 @@ val map_preserving_phys_equal : 'a t -> f:('a -> 'a) -> 'a t
 (* val folding_map : 'a t -> init:'b -> f:('b -> 'a -> 'b * 'c) -> 'c t *)
 (* val folding_mapi : 'a t -> init:'b -> f:(int -> 'b -> 'a -> 'b * 'c) ->
    'c t *)
+
+val fold_map_until :
+     'a t
+  -> init:'accum
+  -> f:('accum -> 'a -> ('accum * 'b, 'final) Continue_or_stop.t)
+  -> finish:('accum * 'b t -> 'final)
+  -> 'final
+
 val fold_map : 'a t -> init:'b -> f:('b -> 'a -> 'b * 'c) -> 'b * 'c t
 
 (* val fold_mapi :

sledge/src/llair/term.ml

@@ -720,61 +720,45 @@ let rec simp_extract agg off len =
         simp_memory len e
     (* ⟨n,a⟩[0,n) ==> ⟨n,a⟩ *)
     | Ap2 (Memory, n, _) when equal off zero && equal n len -> agg
-    (* (α₀^…^αᵢ^…^αⱼ^…) [0+n₀+…+nᵢ₋₁, nᵢ+…+nⱼ) ==> αᵢ^…^αⱼ where nₓ ≡ |αₓ| *)
-    | ApN (Concat, na1N) ->
-        let n = Vector.length na1N in
-        (* invariant: oI = ∑ᵥ₌₀ⁱ⁻¹ nᵥ *)
-        let rec find_off oI i =
-          [%Trace.call fun {pf} -> pf "o_0^%i = %a" i pp oI]
-          ;
-          ( if i = n then Ap3 (Extract, agg, off, len)
-          else
-            match Vector.get na1N i with
-            | Ap2 (Memory, nI, _) | Ap3 (Extract, _, _, nI) -> (
-              match (oI, off) with
-              | Integer {data= y}, Integer {data= z} when Z.gt y z ->
-                  Ap3 (Extract, agg, off, len)
-              | _ when not (equal oI off) ->
-                  find_off (simp_add2 oI nI) (i + 1)
-              | _ ->
-                  (* invariant: lIJ = ∑ᵥ₌ᵢʲ⁻¹ nᵥ *)
-                  let rec find_len lIJ j =
-                    [%Trace.call fun {pf} -> pf "l_%i^%i = %a" i j pp lIJ]
-                    ;
-                    ( if j = n then find_off (simp_add2 oI nI) (i + 1)
+    (* For (α₀^α₁)[o,l) there are 3 cases:
+     *
+     * ⟨...⟩^⟨...⟩
+     *  [,)
+     * o < o+l ≤ |α₀| : (α₀^α₁)[o,l) ==> α₀[o,l) ^ α₁[0,0)
+     *
+     * ⟨...⟩^⟨...⟩
+     *   [  ,  )
+     * o ≤ |α₀| < o+l : (α₀^α₁)[o,l) ==> α₀[o,|α₀|-o) ^ α₁[0,l-(|α₀|-o))
+     *
+     * ⟨...⟩^⟨...⟩
+     *        [,)
+     * |α₀| ≤ o : (α₀^α₁)[o,l) ==> α₀[o,0) ^ α₁[o-|α₀|,l)
+     *
+     * So in general:
+     *
+     * (α₀^α₁)[o,l) ==> α₀[o,l₀) ^ α₁[o₁,l-l₀)
+     * where l₀ = max 0 (min l |α₀|-o)
+     *       o₁ = max 0 o-|α₀|
+     *)
+    | ApN (Concat, na1N) -> (
+      match len with
+      | Integer {data= l} ->
+          Vector.fold_map_until na1N ~init:(l, off)
+            ~f:(fun (l, oI) naI ->
+              let nI = agg_size_exn naI in
+              if Z.equal Z.zero l then
+                Continue ((l, oI), simp_extract naI oI zero)
               else
-                      match Vector.get na1N j with
-                      | Ap2 (Memory, nJ, _) | Ap3 (Extract, _, _, nJ) -> (
-                          let lIJ = simp_add2 lIJ nJ in
-                          match (lIJ, len) with
-                          (* (α₀^…^αᵢ^…) [0+n₀+…+nᵢ₋₁, l) ==> (αᵢ^…)[0,l)
-                             where nₓ ≡ |αₓ| *)
-                          | Integer {data= y}, Integer {data= z}
-                            when Z.gt y z ->
-                              let naIN =
-                                Vector.sub ~pos:i ~len:(n - i) na1N
-                              in
-                              simp_extract (simp_concat naIN) zero len
-                          | Integer {data= y}, Integer {data= z}
-                            when Z.gt y z ->
-                              Ap3 (Extract, agg, off, len)
-                          | _ when not (equal lIJ len) ->
-                              find_len lIJ (j + 1)
-                          | _ ->
-                              let naIJ =
-                                Vector.sub ~pos:i ~len:(j - i + 1) na1N
-                              in
-                              simp_concat naIJ )
-                      | _ -> violates invariant agg )
-                    |>
-                    [%Trace.retn fun {pf} -> pf "%a" pp]
-                  in
-                  find_len zero i )
-            | _ -> violates invariant agg )
-          |>
-          [%Trace.retn fun {pf} -> pf "%a" pp]
-        in
-        find_off zero 0
+                let oI_nI = simp_sub oI nI in
+                match oI_nI with
+                | Integer {data} ->
+                    let oJ = if Z.sign data <= 0 then zero else oI_nI in
+                    let lI = Z.(max zero (min l (neg data))) in
+                    let l = Z.(l - lI) in
+                    Continue ((l, oJ), simp_extract naI oI (integer lI))
+                | _ -> Stop (Ap3 (Extract, agg, off, len)) )
+            ~finish:(fun (_, e1N) -> simp_concat e1N)
+      | _ -> Ap3 (Extract, agg, off, len) )
     (* α[o,l) *)
     | _ -> Ap3 (Extract, agg, off, len) )
   |>