(* example dimension *)
let n = 5

(* example adjacency matrix *)
(* a : int array array *)
let a = [| [| 1; 0; 0; 1; 0 |];
	   [| 0; 0; 1; 0; 0 |];
	   [| 0; 1; 0; 0; 0 |];
	   [| 0; 0; 1; 1; 1 |];
	   [| 0; 0; 1; 0; 1 |] |]

(* input program *)
(* f : int array array -> int array -> int array *)
let f a v =
  let v' = Array.make n 0 in
  for i = 0 to n-1 do
    for j = 0 to n-1 do
      v'.(i) <- v'.(i) + a.(i).(j) * v.(j)
    done
  done;
  v'

(* example vector *)
let v = [| 3; 1; 5; -2; 4 |]

(* example execution *)
let _ = assert (f a v = [|1; 5; 1; 7; 9|])

(* hack for handling semicolon properly (due to Ken Shan) *)

(* for_to : int -> int -> (int -> 'a -> 'a) -> 'a -> 'a *)
let rec for_to n p f init =
    if n > p then init else for_to n (p-1) f (f p init)

(* id : 'a -> 'a *)
let id c = c

(* semi : 'a -> 'b -> 'b *)
let semi c1 c2 = c1; c2

(* solution 1 ----------------------------------------------------- *)

(* rewriting *)

(* f1 : int array array -> int array -> int array *)
let f1 a = fun v ->
  let v' = Array.make n 0 in
  for_to 0 (n-1) (fun i ->
    for_to 0 (n-1) (fun j ->
      if a.(i).(j) = 1 then
        semi (v'.(i) <- v'.(i) + v.(j))
      else id))
  v'

(* example execution *)
let _ = assert (f1 a v = f a v)

(* staging *)

(* semi' : ('a, 'b) code -> ('a, 'c) code -> ('a, 'c) code *)
let semi' c1 c2 = .<(.~c1; .~c2)>.

(* f1' : int array array -> ('a, int array -> int array) code *)
let f1' a = .<fun v ->
  let v' = Array.make n 0 in .~(
  for_to 0 (n-1) (fun i ->
    for_to 0 (n-1) (fun j ->
      if a.(i).(j) = 1 then
	semi' .< v'.(i) <- v'.(i) + v.(j) >.
      else id))
  .<v'>.)>.

(* example execution *)
let _ = f1' a

(*
# f1' a;;
- : ('a, int array -> int array) code =
.<fun v ->
   let v' = Array.make 5 0 in
   v'.(0) <- v'.(0) + v.(0);
   v'.(0) <- v'.(0) + v.(3);
   v'.(1) <- v'.(1) + v.(2);
   v'.(2) <- v'.(2) + v.(1);
   v'.(3) <- v'.(3) + v.(2);
   v'.(3) <- v'.(3) + v.(3);
   v'.(3) <- v'.(3) + v.(4);
   v'.(4) <- v'.(4) + v.(2);
   v'.(4) <- v'.(4) + v.(4);
   v'>.
*)

(* example execution *)
let _ = assert ((.! (f1' a)) v = f a v)

(* solution 2 ----------------------------------------------------- *)

(* rewriting *)

(* non_zeros : int array -> int *)
let rec non_zeros v =
  let n = ref 0 in
  for i = 0 to Array.length v - 1 do
    if v.(i) <> 0 then n := !n + 1
  done;
  !n

(* example threshold *)
let threshold = 3

(* f2 : int array array -> int array -> int array *)
let f2 a = fun v ->
  let v' = Array.make n 0 in
  for_to 0 (n-1) (fun i ->
    if non_zeros a.(i) < threshold
    then for_to 0 (n-1) (fun j ->
	   if a.(i).(j) = 1 then
	     semi (v'.(i) <- v'.(i) + v.(j))
	   else id)
    else semi
	 (for j = 0 to (let n' = n-1 in n') do
	    v'.(i) <- v'.(i) + a.(i).(j) * v.(j)
	  done))
  v'

(* example execution *)
let _ = assert (f2 a v = f a v)

(* staging *)

(* f2' : int array array -> ('a, int array -> int array) code *)
let f2' a = .<fun v ->
  let v' = Array.make n 0 in .~(
  for_to 0 (n-1) (fun i ->
    if non_zeros a.(i) < threshold
    then for_to 0 (n-1) (fun j ->
	   if a.(i).(j) = 1 then
	     semi' .< v'.(i) <- v'.(i) + v.(j) >.
	   else id)
    else semi'
      .< for j = 0 to .~(let n' = n-1 in .<n'>.) do
	   v'.(i) <- v'.(i) + a.(i).(j) * v.(j)
	 done >.)
  .<v'>.)>.

(* example execution *)
let _ = assert ((.! (f2' a)) v = f a v)

(*
# f2' a;;
- : ('a, int array -> int array) code =
.<fun v ->
   let v' = Array.make 5 0 in
   v'.(0) <- v'.(0) + v.(0);
   v'.(0) <- v'.(0) + v.(3);
   v'.(1) <- v'.(1) + v.(2);
   v'.(2) <- v'.(2) + v.(1);
   for j = 0 to 4 do
     v'.(3) <- v'.(3) + a.(3).(j) * v.(j)
   done;
   v'.(4) <- v'.(4) + v.(2);
   v'.(4) <- v'.(4) + v.(4);
   v'>.
*)