Append source region for AL algorithms (#107)

spectec/src/al/al_util.ml

@@ -200,15 +200,15 @@ let unwrap_cate e =
   | CatE (e1, e2) -> e1, e2
   | _ -> fail_expr "unwrap_cate" e
 
-let name_of_algo = function
+let name_of_algo algo = match algo.it with
   | RuleA (name, _, _) -> Print.string_of_atom name
   | FuncA (name, _, _) -> name
 
-let params_of_algo = function
+let params_of_algo algo = match algo.it with
   | RuleA (_, params, _) -> params
   | FuncA (_, params, _) -> params
 
-let body_of_algo = function
+let body_of_algo algo = match algo.it with
   | RuleA (_, _, body) -> body
   | FuncA (_, _, body) -> body
 

spectec/src/al/ast.ml

@@ -152,6 +152,7 @@ and instr' =
 
 (* Algorithms *)
 
-type algorithm =                           (* `algorithm` x`(`expr*`)` `{`instr*`}` *)
+type algorithm = algorithm' phrase
+and algorithm' =                          (* `algorithm` f`(`expr*`)` `{`instr*`}` *)
   | RuleA of atom * expr list * instr list (* reduction rule *)
   | FuncA of id * expr list * instr list   (* helper function *)

spectec/src/al/eq.ml

@@ -167,7 +167,7 @@ and eq_instrs il1 il2 =
 
 
 let eq_algos al1 al2 =
-  match al1, al2 with
+  match al1.it, al2.it with
   | RuleA (a1, el1, il1), RuleA (a2, el2, il2) ->
     a1 = a2 &&
     eq_exprs el1 el2 &&

spectec/src/al/print.ml

@@ -352,7 +352,7 @@ let string_of_instr instr =
   string_of_instr' 0 instr
 let string_of_instrs = string_of_instrs' 0
 
-let string_of_algorithm = function
+let string_of_algorithm algo = match algo.it with
   | RuleA (a, params, instrs) ->
     "execution_of_" ^ string_of_atom a
     ^ List.fold_left
@@ -617,7 +617,7 @@ and structured_string_of_instrs' depth instrs =
 let structured_string_of_instr = structured_string_of_instr' 0
 let structured_string_of_instrs = structured_string_of_instrs' 0
 
-let structured_string_of_algorithm = function
+let structured_string_of_algorithm algo = match algo.it with
   | RuleA (a, params, instrs) ->
       "execution_of_" ^ string_of_atom a
       ^ List.fold_left

spectec/src/al/walk.ml

@@ -134,6 +134,7 @@ let rec walk_instr f (instr:instr) : instr list =
 
 and walk_instrs f = walk_instr f |> List.concat_map
 
-let walk f algo = match algo with
+let walk' f algo' = match algo' with
   | RuleA (a, params, body) -> RuleA (a, params, walk_instrs f body)
   | FuncA (id, params, body) -> FuncA (id, params, walk_instrs f body)
+let walk f algo = Source.map (walk' f) algo

spectec/src/backend-interpreter/ds.ml

@@ -24,7 +24,7 @@ let func_map: func_map ref = ref Map.empty
 let to_map algos =
   let f acc algo =
     let rmap, fmap = acc in
-    match algo with
+    match algo.it with
     | RuleA (atom, _, _) ->
         let name = Print.string_of_atom atom in
         Map.add name algo rmap, fmap

spectec/src/backend-interpreter/interpreter.ml

@@ -665,8 +665,15 @@ and create_context (name: string) (args: value list) : AlContext.mode =
   let params = params_of_algo algo in
   let body = body_of_algo algo in
 
-  if List.length args <> List.length params then
-    raise (Exception.InvalidArg ("Args number mismatch for algorithm " ^ name));
+  if List.length args <> List.length params then (
+    error
+      algo.at
+      (Printf.sprintf "Expected %d arguments for the algorithm `%s` but %d arguments are given"
+        (List.length params)
+        name
+        (List.length args))
+      (string_of_value (CaseV (name, args)))
+  );
 
   let env =
     Env.empty

spectec/src/backend-prose/gen.ml

@@ -146,7 +146,7 @@ let gen_validation_prose il =
 let gen_execution_prose =
   List.map
     (fun algo ->
-      let handle_state = match algo with
+      let handle_state = match algo.it with
       | Al.Ast.RuleA _ -> Il2al.Transpile.insert_state_binding
       | Al.Ast.FuncA _ -> Il2al.Transpile.remove_state
       in

spectec/src/backend-prose/render.ml

@@ -634,7 +634,7 @@ let render_func env fname params instrs =
 let render_def env = function
   | Pred (name, params, instrs) ->
     "\n" ^ render_pred env name params instrs ^ "\n\n"
-  | Algo algo -> (match algo with
+  | Algo algo -> (match algo.it with
     | Al.Ast.RuleA (name, params, instrs) ->
       "\n" ^ render_rule env name params instrs ^ "\n\n"
     | Al.Ast.FuncA (name, params, instrs) ->

spectec/src/backend-splice/splice.ml

@@ -116,12 +116,12 @@ let env_prose env prose =
     let relation = Map.find valid_id env.rel_prose in
     let ralgos = (normalize_id id, prose, ref 0) :: relation.ralgos in
     env.rel_prose <- Map.add valid_id {ralgos} env.rel_prose
-  | Algo (Al.Ast.RuleA ((id, typ), _, _)) ->
+  | Algo ({ it = Al.Ast.RuleA ((id, typ), _, _); _ }) ->
     let id = El.Atom.to_string (id $$ (no_region, ref typ)) in
     let relation = Map.find exec_id env.rel_prose in
     let ralgos = (normalize_id id, prose, ref 0) :: relation.ralgos in
     env.rel_prose <- Map.add exec_id {ralgos} env.rel_prose
-  | Algo (Al.Ast.FuncA (id, _, _)) ->
+  | Algo ({ it = Al.Ast.FuncA (id, _, _); _}) ->
     env.def_prose <- Map.add id {falgo = prose; use = ref 0} env.def_prose
 
 let env (config : config) pdsts odsts elab el pr : env =

spectec/src/il2al/il2il.ml

@@ -4,9 +4,10 @@ Functions that transform IL into IL.
 
 open Il.Ast
 open Il.Eq
+open Util
 open Util.Source
 
-type rgroup = (exp * exp * (prem list)) list
+type rgroup = (exp * exp * (prem list)) phrase list
 
 (* Helpers *)
 
@@ -261,17 +262,18 @@ let apply_template_to_rgroup template (lhs, rhs, prems) =
 
 let unify_lhs' rgroup =
   init_unified_idx();
-  let lhs_group = List.map (function (lhs, _, _) -> lhs) rgroup in
+  let fst = fun (x, _, _) -> x in
+  let lhs_group = List.map (fun r -> fst r.it) rgroup in
   let hd = List.hd lhs_group in
   let tl = List.tl lhs_group in
   let template = List.fold_left overlap hd tl in
-  List.map (apply_template_to_rgroup template) rgroup
+  List.map (Source.map (apply_template_to_rgroup template)) rgroup
 
 let unify_lhs (rname, rgroup) =
   let to_left_assoc (lhs, rhs, prems) = to_left_assoc_cat lhs, rhs, prems in
   let to_right_assoc (lhs, rhs, prems) = to_right_assoc_cat lhs, rhs, prems in
   (* typical f^-1 ∘ g ∘ f *)
-  rname, (rgroup |> List.map to_left_assoc |> unify_lhs' |> List.map to_right_assoc)
+  rname, (rgroup |> List.map (Source.map to_left_assoc) |> unify_lhs' |> List.map (Source.map to_right_assoc))
 
 let apply_template_to_def template def =
   match def.it with

spectec/src/il2al/il2il.mli

@@ -1,6 +1,6 @@
 open Il.Ast
 
-type rgroup = (exp * exp * (prem list)) list
+type rgroup = (exp * exp * (prem list)) Util.Source.phrase list
 
 val transform_expr : (exp -> exp) -> exp -> exp
 val is_unified_id : string -> bool

spectec/src/il2al/manual.ml

@@ -19,7 +19,7 @@ let eval_expr =
       popI result;
       returnI (Some (listE [ result ]))
     ]
-  )
+  ) $ no_region
 
 (* Helper for the manual array_new.data algorithm *)
 
@@ -48,7 +48,7 @@ let group_bytes_by =
       );
       returnI (Some (listE []));
     ]
-  )
+  ) $ no_region
 
 let array_new_data =
   let i32 = caseE (atom_of_name "I32" "numtype", []) in
@@ -118,7 +118,7 @@ let array_new_data =
         []
       );
     ]
-  )
+  ) $ no_region
 
 let manual_algos = [eval_expr; group_bytes_by; array_new_data;]
 

spectec/src/il2al/translate.ml

@@ -74,7 +74,7 @@ let get_params winstr =
     (sprintf "cannot get params of wasm instruction `%s`" (Il.Print.string_of_exp winstr))
 
 let lhs_of_rgroup rgroup =
-  let (lhs, _, _) = List.hd rgroup in
+  let (lhs, _, _) = (List.hd rgroup).it in
   lhs
 
 let name_of_rule rule =
@@ -785,7 +785,7 @@ let translate_helper partial_funcs def =
       |> (if List.mem id partial_funcs then Fun.id else Transpile.ensure_return)
       |> Transpile.flatten_if in
 
-    Some (FuncA (name, params, body))
+    Some (FuncA (name, params, body) $ def.at)
   | _ -> None
 
 
@@ -818,7 +818,7 @@ let in_same_context (lhs1, _, _) (lhs2, _, _) =
 
 let group_contexts xs =
   List.fold_left (fun acc x ->
-    let g1, g2 = List.partition (fun g -> in_same_context (List.hd g) x) acc in
+    let g1, g2 = List.partition (fun g -> in_same_context (List.hd g).it x.it) acc in
     match g1 with
     | [] -> [ x ] :: acc
     | [ g ] -> (x :: g) :: g2
@@ -856,7 +856,7 @@ let insert_deferred = function
 
 (* `reduction` -> `instr list` *)
 let translate_reduction deferred reduction =
-  let _, rhs, prems = reduction in
+  let _, rhs, prems = reduction.it in
 
   (* Translate rhs *)
   translate_rhs rhs
@@ -951,7 +951,7 @@ let translate_context_rgroup lhss sub_algos inner_params =
   let instr_popall = popallI e_vals in
   let instrs_context =
     List.fold_right2 (fun lhs algo acc ->
-      match algo with
+      match algo.it with
       | RuleA (_, params, body) ->
         (* Assume that each sub-algorithms are produced by translate_context,
            i.e., they will always contain instr_popall as their first instruction. *)
@@ -1006,7 +1006,7 @@ let rec translate_rgroup' context winstr instr_name rgroup =
       let instrs' =
         match rgroup |> Util.Lib.List.split_last with
         (* Either case: No premise for the last reduction rule *)
-        | hds, (_, rhs, []) when List.length hds > 0 ->
+        | hds, { it = (_, rhs, []); _ } when List.length hds > 0 ->
           assert (defer_opt = None);
           let blocks = List.map (translate_reduction None) hds in
           let body1 = Transpile.merge_blocks blocks in
@@ -1029,7 +1029,7 @@ let rec translate_rgroup' context winstr instr_name rgroup =
       (try
       let unified_sub_groups =
         rgroup
-        |> List.map un_unify
+        |> List.map (Source.map un_unify)
         |> group_contexts
         |> List.map (fun g -> Il2il.unify_lhs (instr_name, g)) in
 
@@ -1051,7 +1051,7 @@ and get_lhs_stack (exp: Il.exp): Il.exp list =
   else to_exp_list exp
 
 and translate_rgroup (instr_name, rgroup) =
-  let lhs, _, _ = List.hd rgroup in
+  let lhs, _, _ = (List.hd rgroup).it in
   (* TODO: Generalize getting current frame *)
   let lhs_stack = get_lhs_stack lhs in
   let context, winstr = split_lhs_stack instr_name lhs_stack in
@@ -1086,14 +1086,18 @@ and translate_rgroup (instr_name, rgroup) =
     |> Transpile.infer_assert
     |> Transpile.flatten_if
   in
-  RuleA (name, al_params', body)
+
+  let at = rgroup
+    |> List.map at
+    |> over_region in
+  RuleA (name, al_params', body) $ at
 
 
 let rule_to_tup rule =
   match rule.it with
   | Il.RuleD (_, _, _, exp, prems) ->
     match exp.it with
-    | Il.TupE [ lhs; rhs ] -> lhs, rhs, prems
+    | Il.TupE [ lhs; rhs ] -> (lhs, rhs, prems) $ rule.at
     | _ -> error_exp exp "form of reduction rule"
 
 

spectec/src/il2al/transpile.ml

@@ -512,16 +512,19 @@ let remove_state algo =
       }
   in
 
-  match Walk.walk walk_config algo with
-  | FuncA (name, args, body) ->
-    let args' =
-      args
-      |> Lib.List.filter_not is_state
-      |> Lib.List.filter_not is_store
-      |> Lib.List.filter_not is_frame
-    in
-    FuncA (name, args', body)
-  | rule -> rule
+  let algo' = Walk.walk walk_config algo in
+  { algo' with it =
+    match algo'.it with
+    | FuncA (name, args, body) ->
+      let args' =
+        args
+        |> Lib.List.filter_not is_state
+        |> Lib.List.filter_not is_store
+        |> Lib.List.filter_not is_frame
+      in
+      FuncA (name, args', body)
+    | rule -> rule
+  }
 
 let insert_state_binding algo =
   let state_count = ref 0 in
@@ -540,14 +543,17 @@ let insert_state_binding algo =
     }
   in
 
-  match Walk.walk walk_config algo with
-  | FuncA (name, params, body) when !state_count > 0 ->
-    let body = (letI (varE "z", getCurStateE ())) :: body in
-    FuncA (name, params, body)
-  | RuleA (name, params, body) when !state_count > 0 ->
-    let body = (letI (varE "z", getCurStateE ())) :: body in
-    RuleA (name, params, body)
-  | _ -> algo
+  let algo' = Walk.walk walk_config algo in
+  { algo' with it =
+    match algo'.it with
+    | FuncA (name, params, body) when !state_count > 0 ->
+      let body = (letI (varE "z", getCurStateE ())) :: body in
+      FuncA (name, params, body)
+    | RuleA (name, params, body) when !state_count > 0 ->
+      let body = (letI (varE "z", getCurStateE ())) :: body in
+      RuleA (name, params, body)
+    | a -> a
+  }
 
 let insert_frame_binding instrs =
   let frame_count = ref 0 in
@@ -691,8 +697,7 @@ let remove_enter algo =
     | _ -> instr
   in
 
-  let remove_enter' algo =
-    match algo with
+  let remove_enter' = Source.map (function
     | FuncA (name, params, body) ->
         let walk_config =
           {
@@ -711,7 +716,7 @@ let remove_enter algo =
         in
         let body = Walk.walk_instrs walk_config body in
         RuleA (name, params, body)
-  in
+  ) in
 
   let algo' = remove_enter' algo in
   if Eq.eq_algos algo algo' then algo else remove_enter' algo'

spectec/src/util/source.ml

@@ -43,3 +43,6 @@ let (%) at note = (at, note)
 let it {it; _} = it
 let at {at; _} = at
 let note {note; _} = note
+
+(* Utils *)
+let map f {it; at; note} = {it = f it; at; note}

spectec/src/util/source.mli

@@ -25,3 +25,6 @@ val (%) : region -> 'b -> region * 'b
 val it : ('a, 'b) note_phrase -> 'a
 val at : ('a, 'b) note_phrase -> region
 val note : ('a, 'b) note_phrase -> 'b
+
+(* Utils *)
+val map : ('a -> 'a) -> ('a, 'b) note_phrase -> ('a, 'b) note_phrase