Bug fix in parallel_scan

lib/task.ml

@@ -203,36 +203,41 @@ let parallel_for ?(chunk_size=0) ~start ~finish ~body pool =
 
 let parallel_scan pool op elements =
   let pd = get_pool_data pool in
+  let n = Array.length elements in
+  let p = min (n - 1) ((Array.length pd.domains) + 1) in
+  let prefix_s = Array.copy elements in
   let scan_part op elements prefix_sum start finish =
     assert (Array.length elements > (finish - start));
     for i = (start + 1) to finish do
       prefix_sum.(i) <- op prefix_sum.(i - 1) elements.(i)
     done
   in
+  if p < 2 then begin
+    (* Do a sequential scan when number of domains or array's length is less
+    than 2 *)
+    scan_part op elements prefix_s 0 (n - 1);
+    prefix_s
+  end
+  else begin
   let add_offset op prefix_sum offset start finish =
     assert (Array.length prefix_sum > (finish - start));
     for i = start to finish do
       prefix_sum.(i) <- op offset prefix_sum.(i)
     done
   in
-  let n = Array.length elements in
-  let p = (Array.length pd.domains) + 1 in
-  let prefix_s = Array.copy elements in
 
   parallel_for pool ~chunk_size:1 ~start:0 ~finish:(p - 1)
   ~body:(fun i ->
     let s = (i * n) / (p ) in
     let e = (i + 1) * n / (p ) - 1 in
     scan_part op elements prefix_s s e);
 
-  if (p > 2) then begin
   let x = ref prefix_s.(n/p - 1) in
   for i = 2 to p do
       let ind = i * n / p - 1 in
       x := op prefix_s.(ind) !x;
       prefix_s.(ind) <- !x
-  done
-  end;
+  done;
 
   parallel_for pool ~chunk_size:1 ~start:1 ~finish:(p - 1)
   ~body:( fun i ->
@@ -243,3 +248,4 @@ let parallel_scan pool op elements =
     );
 
   prefix_s
+  end

test/dune

@@ -92,3 +92,9 @@
  (libraries domainslib)
  (modules backtrace)
  (modes native))
+
+(test
+ (name off_by_one)
+ (libraries domainslib)
+ (modules off_by_one)
+ (modes native))

test/off_by_one.ml

@@ -0,0 +1,21 @@
+open Domainslib
+
+let print_array a =
+  let b = Buffer.create 25 in
+  Buffer.add_string b "[|";
+  Array.iter (fun elem -> Buffer.add_string b (string_of_int elem ^ "; ")) a;
+  Buffer.add_string b "|]";
+  Buffer.contents b
+
+let r = Array.init 20 (fun i -> i + 1)
+
+let scan_task num_doms =
+  let pool = Task.setup_pool ~num_additional_domains:num_doms () in
+  let a = Task.run pool (fun () -> Task.parallel_scan pool (+) (Array.make 20 1)) in
+  Task.teardown_pool pool;
+  Printf.printf "%i:  %s\n%!" num_doms (print_array a);
+  assert (a = r)
+;;
+for num_dom=0 to 21 do
+  scan_task num_dom;
+done