Optimize non zero-preserving Vector.map and Vector.map2i

src/cdomains/affineEquality/sparseImplementation/sparseVector.ml

@@ -96,7 +96,19 @@ module SparseVector: AbstractVector =
         {entries = entries'; len = v.len + 1}
 
     let map f v = 
+      (*
       of_list (List.map f (to_list v))
+      *)
+      let f_zero = f A.zero in
+      let rec map2_helper acc vec i =
+        match vec with
+        | [] when i >= v.len || f_zero =: A.zero -> List.rev acc
+        | []  -> map2_helper ((i, f_zero) :: acc) [] (i  + 1)
+        | (idx, value) :: xs when idx = i -> let new_val = f value in if new_val <>: A.zero then map2_helper ((idx, new_val) :: acc) xs (i + 1) else map2_helper acc xs (i + 1)
+        | (idx, _) :: xs when idx > i -> if f_zero <>: A.zero then map2_helper ((i, f_zero) :: acc) vec (i + 1) else map2_helper acc vec (i + 1)
+        | (_, _) :: _ -> failwith "This should not happen"
+      in
+      {entries = map2_helper [] v.entries 0; len = v.len}
 
     let map_f_preserves_zero f v = (* map for functions f such that f 0 = 0 since f won't be applied to zero values. See also map *)
       let entries' = List.filter_map (
@@ -132,8 +144,29 @@ module SparseVector: AbstractVector =
 
     let map2_f_preserves_zero f v1 v2 = Timing.wrap "map2_f_preserves_zero" (map2_f_preserves_zero f v1) v2
 
-    let map2i f v v' = (* TODO: optimize! *)
-      of_list (List.map2i f (to_list v) (to_list v'))
+    let map2i f v1 v2 = (* TODO: optimize! *)
+      if v1.len <> v2.len then raise (Invalid_argument "Unequal lengths") else
+        (*of_list (List.map2i f (to_list v) (to_list v'))*)
+        let f_rem_zero idx acc e1 e2 =
+          let r = f idx e1 e2 in 
+          if r =: A.zero then acc else (idx, r)::acc
+        in  
+        let rec aux acc vec1 vec2 i =
+          match vec1, vec2 with 
+          | [], [] when i = v1.len -> acc 
+          | [], [] -> aux (f_rem_zero i acc A.zero A.zero) [] [] (i + 1)
+          | [], (yidx, yval)::ys when i = yidx -> aux (f_rem_zero i acc A.zero yval) [] ys (i + 1)
+          | (xidx, xval)::xs, [] when i = xidx -> aux (f_rem_zero i acc xval A.zero) xs [] (i + 1)
+          | [], (_, _)::_ | (_, _)::_, [] -> aux (f_rem_zero i acc A.zero A.zero) vec1 vec2 (i + 1) (* When one vec is not zero_vec, but has implicit zeroes at front *)
+          | (xidx, xval)::xs, (yidx, yval)::ys -> 
+            if xidx <> i && yidx <> i then aux (f_rem_zero i acc A.zero A.zero) vec1 vec2 (i+1) (* When both vectors have implicit zeroes at front *)
+            else
+            match xidx - yidx with (* Here at least one of the idx is i, which is the smaller one *)
+            | d when d < 0 -> aux (f_rem_zero i acc xval A.zero) xs vec2 (i + 1)
+            | d when d > 0 -> aux (f_rem_zero i acc A.zero yval) vec1 ys (i + 1)
+            | _            -> aux (f_rem_zero i acc xval yval) xs ys (i + 1)
+        in
+        {entries = List.rev (aux [] v1.entries v2.entries 0); len = v1.len}
 
     let map2i_f_preserves_zero f v v' = failwith "TODO"
 

tests/unit/cdomains/affineEqualityDomain/sparseImplementation/sparseMatrixImplementationTest.ml

@@ -8,9 +8,12 @@ open SparseVector
 open ListMatrix
 open ArrayVector
 open ArrayMatrix
+open ConvenienceOps
+
 module D = SharedFunctions.Mpqf
 module Vector = SparseVector (D)
 module Matrix = ListMatrix (D) (SparseVector)
+include ConvenienceOps(D)
 
 (** Shorthands for common functions. *)
 let int x = D.of_int x
@@ -303,25 +306,65 @@ let int_domain_to_rational _ =
   in
   normalize_and_assert int_matrix normalized_matrix
 
-let tests =
-  "SparseMatrixImplementationTest"
-  >::: [
-    "can solve a standard normalization" >:: standard_normalize;
-    "does sort already reduzed" >:: does_just_sort;
-    "does eliminate dependent rows" >:: does_eliminate_dependent_rows;
-    "can handle float domain" >:: does_handle_floats;
-    "can handle fraction domain" >:: does_handle_fractions;
-    "does negate negative matrix" >:: does_negate_negative;
-    "does not change already normalized matrix" >:: does_not_change_normalized_matrix;
-    "m1 is covered by m2" >:: is_covered_by_simple;
-    "m1 is covered by m2 with vector in first row" >:: is_covered_by_vector_first_row;
-    "zero vector is covered by m2" >:: is_zero_vec_covered;
-    "m1 is not covered by m2" >:: is_not_covered;
-    "m1 is covered by m2 with big matrix" >:: is_covered_big;
-    "does not change an empty matrix" >:: normalize_empty;
-    "can correctly normalize a two column matrix" >:: normalize_two_columns;
-    "can handle a rational solution" >:: int_domain_to_rational;
-    "m1 is covered by m2 with big matrix2" >:: is_covered_big2;
-  ]
+
+let vectorMap2i _ =
+  let v1 = Vector.of_list [int 0; int 1; int 0; int 2; int 3; int 0; int 4; int 0; int 1] in
+  let v2 = Vector.of_list [int 4; int 0; int 0; int 0; int 5; int 6; int 0; int 0; int 2] in
+  let result = Vector.map2i (fun i x y -> (int i) *: (x +: y)) v1 v2 in
+  let expected = Vector.of_list [int 0; int 1; int 0; int 6; int 32; int 30; int 24; int 0; int 24] in
+  assert_equal expected result
+
+
+let vectorMap2i_empty _ = 
+  let v1 = Vector.of_list [] in
+  let v2 = Vector.of_list [] in
+  let result = Vector.map2i (fun i x y -> (int i) *: (x +: y)) v1 v2 in
+  let expected = Vector.of_list [] in
+  assert_equal expected result
+
+let vectorMap2i_one_zero _ = 
+    let v1 = Vector.of_list [int 0; int 0; int 0; int 0] in
+    let v2 = Vector.of_list [int 1; int 2; int 3; int 4] in
+    let result = Vector.map2i (fun i x y -> (int i) *: (x +: y)) v1 v2 in
+    let expected = Vector.of_list [int 0; int 2; int 6; int 12] in
+    assert_equal expected result
+
+let vectorMap _ =
+  let v1 = Vector.of_list [int 0; int 1; int 2; int 0; int 0; int 3; int 4; int 0; int 0; int 5] in
+  let result = Vector.map (fun x -> x +: int 1) v1 in
+  let expected = Vector.of_list [int 1; int 2; int 3; int 1; int 1; int 4; int 5; int 1; int 1; int 6] in
+  assert_equal expected result
+
+let vectorMap_zero_preserving_normal _ =
+  let v1 = Vector.of_list [int 0; int 1; int 2; int 0; int 0; int 4; int 5; int 0; int 0;] in
+  let result = Vector.map (fun x -> x *: x) v1 in
+  let expected = Vector.of_list [int 0; int 1; int 4; int 0; int 0; int 16; int 25; int 0; int 0;] in
+  assert_equal expected result
+
+  let tests =
+    "SparseMatrixImplementationTest"
+    >::: [
+      "can solve a standard normalization" >:: standard_normalize;
+      "does sort already reduzed" >:: does_just_sort;
+      "does eliminate dependent rows" >:: does_eliminate_dependent_rows;
+      "can handle float domain" >:: does_handle_floats;
+      "can handle fraction domain" >:: does_handle_fractions;
+      "does negate negative matrix" >:: does_negate_negative;
+      "does not change already normalized matrix" >:: does_not_change_normalized_matrix;
+      "m1 is covered by m2" >:: is_covered_by_simple;
+      "m1 is covered by m2 with vector in first row" >:: is_covered_by_vector_first_row;
+      "zero vector is covered by m2" >:: is_zero_vec_covered;
+      "m1 is not covered by m2" >:: is_not_covered;
+      "m1 is covered by m2 with big matrix" >:: is_covered_big;
+      "does not change an empty matrix" >:: normalize_empty;
+      "can correctly normalize a two column matrix" >:: normalize_two_columns;
+      "can handle a rational solution" >:: int_domain_to_rational;
+      "m1 is covered by m2 with big matrix2" >:: is_covered_big2;
+      "map2i two vectors" >:: vectorMap2i;
+      "map2i two empty vectors" >:: vectorMap2i_empty;
+      "map2i one zero vector" >:: vectorMap2i_one_zero;
+      "map one vector" >:: vectorMap;
+      "map zero preserving normal" >:: vectorMap_zero_preserving_normal;
+    ]
 
 let () = run_test_tt_main tests