Sibir/add tests

EML/lib/backend/ricsv/peephole.ml

@@ -248,12 +248,11 @@ let same_memory_key (base1, offset1) (base2, offset2) =
   equal_reg base1 base2 && offset1 = offset2
 ;;
 
-let rec find_cached_load key = function
-  | [] -> None
-  | (cached_key, cached_register) :: rest ->
-    if same_memory_key cached_key key
-    then Some cached_register
-    else find_cached_load key rest
+let find_cached_load key cache =
+  List.find_map
+    (fun (cached_key, cached_register) ->
+       if same_memory_key cached_key key then Some cached_register else None)
+    cache
 ;;
 
 let remove_cached_key key cache =
@@ -326,24 +325,26 @@ let can_prove_store_dead ~allow_drop_at_block_end slot following_instructions =
   let rec walk = function
     | [] -> allow_drop_at_block_end
     | instruction :: rest ->
-      if reads_slot slot instruction
-      then false
-      else if stores_slot slot instruction
-      then true
-      else if writes_slot_base slot instruction
-      then false
-      else walk rest
+      (match
+         ( reads_slot slot instruction
+         , stores_slot slot instruction
+         , writes_slot_base slot instruction )
+       with
+       | true, _, _ -> false
+       | _, true, _ -> true
+       | _, _, true -> false
+       | _ -> walk rest)
   in
   walk following_instructions
 ;;
 
 let eliminate_dead_stores_in_block ~allow_drop_at_block_end block =
   let rec loop changed acc = function
     | [] -> List.rev acc, changed
-    | (Sd (_, slot) as store_instruction) :: rest ->
-      if can_prove_store_dead ~allow_drop_at_block_end slot rest
-      then loop true acc rest
-      else loop changed (store_instruction :: acc) rest
+    | Sd (_, slot) :: rest when can_prove_store_dead ~allow_drop_at_block_end slot rest ->
+      loop true acc rest
+    | (Sd (_, _) as store_instruction) :: rest ->
+      loop changed (store_instruction :: acc) rest
     | instruction :: rest -> loop changed (instruction :: acc) rest
   in
   loop false [] block
@@ -389,20 +390,20 @@ let find_redundant_restore_store loaded_register loaded_slot following_instructi
   in
   let rec search prefix = function
     | [] -> List.rev_append prefix [], false
+    | instruction :: rest
+      when is_barrier instruction || writes_slot_base loaded_slot instruction ->
+      List.rev_append prefix (instruction :: rest), false
+    | Sd (stored_register, store_slot) :: rest
+      when same_memory_key loaded_slot store_slot
+           && equal_reg stored_register loaded_register ->
+      List.rev_append prefix rest, true
+    | Sd (stored_register, store_slot) :: rest when same_memory_key loaded_slot store_slot
+      -> List.rev_append prefix (Sd (stored_register, store_slot) :: rest), false
     | instruction :: rest ->
-      if is_barrier instruction || writes_slot_base loaded_slot instruction
-      then List.rev_append prefix (instruction :: rest), false
-      else (
-        match instruction with
-        | Sd (stored_register, store_slot) when same_memory_key loaded_slot store_slot ->
-          if equal_reg stored_register loaded_register
-          then List.rev_append prefix rest, true
-          else List.rev_append prefix (instruction :: rest), false
-        | _ ->
-          (match write_reg instruction with
-           | Some written_register when equal_reg written_register loaded_register ->
-             List.rev_append prefix (instruction :: rest), false
-           | _ -> search (instruction :: prefix) rest))
+      (match write_reg instruction with
+       | Some written_register when equal_reg written_register loaded_register ->
+         List.rev_append prefix (instruction :: rest), false
+       | _ -> search (instruction :: prefix) rest)
   in
   search [] following_instructions
 ;;

EML/lib/frontend/parser.ml

@@ -82,6 +82,44 @@ let is_operator_char_infix = function
   | '|' | ':' | _ -> false
 ;;
 
+let is_custom_power_op op =
+  String.length op >= 2 && String.equal (String.sub op ~pos:0 ~len:2) "**"
+;;
+
+let first_char op = String.get op 0
+
+let is_custom_mul_op op =
+  (not (is_custom_power_op op))
+  &&
+  match first_char op with
+  | '*' | '/' | '%' -> true
+  | _ -> false
+;;
+
+let is_custom_add_op op =
+  match first_char op with
+  | '+' | '-' -> true
+  | _ -> false
+;;
+
+let is_custom_concat_op op =
+  match first_char op with
+  | '@' | '^' -> true
+  | _ -> false
+;;
+
+let is_custom_cmp_op op =
+  match first_char op with
+  | '=' | '<' | '>' | '|' | '&' | '$' -> true
+  | _ -> false
+;;
+
+let is_custom_lowest_op op =
+  match first_char op with
+  | '!' | '?' | '~' | '.' -> true
+  | _ -> false
+;;
+
 let white_space = take_while Char.is_whitespace
 let token s = white_space *> string s
 let token1 s = white_space *> s
@@ -303,6 +341,15 @@ let parse_expr_bin_oper parse_bin_op tkn =
   token tkn *> return (fun e1 e2 -> ExpBinOper (parse_bin_op, e1, e2))
 ;;
 
+let parse_right_associative expr oper =
+  let rec parse () =
+    expr
+    >>= fun left ->
+    oper >>= (fun combine -> parse () >>| fun right -> combine left right) <|> return left
+  in
+  parse ()
+;;
+
 let multiply = parse_expr_bin_oper Multiply "*"
 let division = parse_expr_bin_oper Division "/"
 let plus = parse_expr_bin_oper Plus "+"
@@ -322,17 +369,14 @@ let compare =
 let and_op = parse_expr_bin_oper And "&&"
 let or_op = parse_expr_bin_oper Or "||"
 
-let parse_custom_infix =
-  white_space *> take_while1 is_operator_char_infix
-  >>| fun op -> fun e1 e2 -> ExpBinOper (Custom op, e1, e2)
-;;
-
-let parse_custom_infix_except except =
+let parse_custom_infix_when pred =
   white_space *> take_while1 is_operator_char_infix
   >>= fun op ->
-  if List.mem except op ~equal:String.equal
+  if Option.is_some (builtin_op_of_string op)
   then fail "builtin"
-  else return (fun e1 e2 -> ExpBinOper (Custom op, e1, e2))
+  else if pred op
+  then return (fun e1 e2 -> ExpBinOper (Custom op, e1, e2))
+  else fail "custom_op_mismatch"
 ;;
 
 let parse_expr_ident = parse_ident >>| fun x -> ExpIdent x
@@ -534,32 +578,35 @@ let parse_top_expr parse_expr =
     ]
 ;;
 
-let lower_than_ops1 = [ "+"; "-"; ">="; "<="; "<>"; "="; ">"; "<"; "&&"; "||" ]
-let lower_than_ops2 = [ ">="; "<="; "<>"; "="; ">"; "<"; "&&"; "||" ]
-let lower_than_cmp = [ "&&"; "||" ]
-
 let parse_exp_apply e right =
   let app = parse_expr_apply e right in
   let app = parse_expr_unar_oper app <|> app in
+  let power = parse_right_associative app (parse_custom_infix_when is_custom_power_op) in
   let ops1 =
     parse_left_associative
-      app
-      (parse_custom_infix_except lower_than_ops1 <|> multiply <|> division)
-      app
+      power
+      (parse_custom_infix_when is_custom_mul_op <|> multiply <|> division)
+      power
   in
   let ops2 =
     parse_left_associative
       ops1
-      (parse_custom_infix_except lower_than_ops2 <|> plus <|> minus)
+      (parse_custom_infix_when is_custom_add_op <|> plus <|> minus)
       ops1
   in
+  let concat =
+    parse_right_associative ops2 (parse_custom_infix_when is_custom_concat_op)
+  in
   let cmp =
     parse_left_associative
-      ops2
-      (parse_custom_infix_except lower_than_cmp <|> compare)
-      ops2
+      concat
+      (parse_custom_infix_when is_custom_cmp_op
+       <|> parse_custom_infix_when is_custom_lowest_op
+       <|> compare)
+      concat
   in
-  parse_left_associative cmp (parse_custom_infix <|> and_op <|> or_op) cmp
+  let bool_and = parse_right_associative cmp and_op in
+  parse_right_associative bool_and or_op
 ;;
 
 let parse_expr =

EML/lib/middleend/inferencer.ml

@@ -284,7 +284,7 @@ module TypeEnv = struct
 
   let apply subst env = Map.map env ~f:(Scheme.apply subst)
   let find = Map.find
-  let keys t = Map.keys t
+  let keys = Map.keys
 
   let initial_env =
     let open Base.Map in
@@ -481,7 +481,7 @@ let get_binop_arg_res env op =
       , Substitution.apply subst arg_ty2
       , Substitution.apply subst res_ty
       , subst )
-  | op ->
+  | _ ->
     let* ty1, ty2, ty_res = infer_binop_type op in
     return (ty1, ty2, ty_res, Substitution.empty)
 ;;

EML/lib/middleend/resolve_builtins.ml

@@ -22,17 +22,14 @@ let names_of_pattern p =
 
 let names_of_bind (pat, _) = names_of_pattern pat
 
-let rec resolve_expr scope e =
-  match e with
+let rec resolve_expr scope = function
   | ExpBinOper (Custom op, e1, e2) ->
     let e1' = resolve_expr scope e1 in
     let e2' = resolve_expr scope e2 in
-    (match List.mem scope op ~equal:String.equal with
-     | true -> ExpBinOper (Custom op, e1', e2')
-     | false ->
-       (match builtin_op_of_string op with
-        | Some b -> ExpBinOper (b, e1', e2')
-        | None -> ExpBinOper (Custom op, e1', e2')))
+    (match List.mem scope op ~equal:String.equal, builtin_op_of_string op with
+     | true, _ -> ExpBinOper (Custom op, e1', e2')
+     | false, Some b -> ExpBinOper (b, e1', e2')
+     | false, None -> ExpBinOper (Custom op, e1', e2'))
   | ExpIdent x -> ExpIdent x
   | ExpConst c -> ExpConst c
   | ExpBranch (c, t, o) ->

EML/tests/additional_tests/custom_op_pipe.ml

@@ -0,0 +1,7 @@
+let ( ~> ) x f = f x
+let succ x = x + 1
+
+let main =
+  let () = print_int (10 ~>succ) in
+  0
+;;

EML/tests/additional_tests/custom_op_right_assoc_diff.ml

@@ -0,0 +1,6 @@
+let ( ** ) x y = x - y
+
+let main =
+  let () = print_int (10 ** 3 ** 1) in
+  0
+;;

EML/tests/additional_tests/custom_op_via_op.ml

@@ -0,0 +1,7 @@
+let ( ** ) x y = x * y
+let ( +++ ) x y = (x ** y) + 1
+
+let main =
+  let () = print_int (3 +++ 4) in
+  0
+;;

EML/tests/anf_tests.ml

@@ -304,3 +304,26 @@ let%expect_test "anf_roundtrip_types_partial" =
    | Error e -> Printf.printf "FAIL: %s\n" e);
   [%expect {| OK: types preserved after ANF round-trip |}]
 ;;
+
+let%expect_test "custom_infix_operator_lowers_to_app" =
+  parse_and_anf "let ( =^.^= ) x y = (x * 10) + y";
+  [%expect
+    {|
+[(AnfValue (NonRec,
+    ("=^.^=", 2,
+     (AnfExpr
+        (ComplexLambda ([(PatVariable "x")],
+           (AnfExpr
+              (ComplexLambda ([(PatVariable "y")],
+                 (AnfLet (NonRec, "anf_t0",
+                    (ComplexBinOper (Multiply, (ImmediateVar "x"),
+                       (ImmediateConst (ConstInt 10)))),
+                    (AnfExpr
+                       (ComplexBinOper (Plus, (ImmediateVar "anf_t0"),
+                          (ImmediateVar "y"))))
+                    ))
+                 )))
+           )))),
+    []))
+  ]|}]
+;;

EML/tests/dune

@@ -26,6 +26,7 @@
   (file ../bin/EML.exe)
   (file Makefile)
   (file ../lib/runtime/rv64_runtime.a)
+  (source_tree additional_tests)
   (source_tree gc_tests)
   (source_tree many_tests)))
 

EML/tests/infer.t

@@ -114,3 +114,16 @@ SPDX-License-Identifier: LGPL-3.0-or-later
   $ make infer many_tests/do_not_type/099.ml 2>&1 | sed -n '1p'
   Inferencer error: Left-hand side error: Only variables are allowed on the left-hand side of let rec.
 
+  $ make infer additional_tests/custom_op_via_op.ml
+  val **: int -> int -> int
+  val +++: int -> int -> int
+  val main: int
+
+  $ make infer additional_tests/custom_op_right_assoc_diff.ml
+  val **: int -> int -> int
+  val main: int
+
+  $ make infer additional_tests/custom_op_pipe.ml
+  val main: int
+  val succ: int -> int
+  val ~>: t0 -> (t0 -> t2) -> t2

EML/tests/inferencer_tests.ml

@@ -440,3 +440,47 @@ let%expect_test "test_ast_pattern_unit_lambda" =
   pretty_printer_infer_simple_expression (ExpLambda (PatUnit, [], ExpConst (ConstInt 1)));
   [%expect {|unit -> int|}]
 ;;
+
+let%expect_test "custom_infix_operator" =
+  pretty_printer_parse_and_infer
+    {| let ( ** ) x y = x * y
+let main = 2 ** 3 |};
+  [%expect
+    {|
+    val **: int -> int -> int
+    val main: int|}]
+;;
+
+let%expect_test "custom_infix_bind_like" =
+  pretty_printer_parse_and_infer
+    {| let ( >>= ) n _ = if n <= 1 then 1 else n * (n - 1)
+let main = 3 >>= 0 |};
+  [%expect
+    {|
+    val >>=: int -> t1 -> int
+    val main: int|}]
+;;
+
+let%expect_test "custom_infix_power" =
+  pretty_printer_parse_and_infer
+    {| let rec ( ^^ ) x n = if n <= 0 then 1 else x * (x ^^ (n - 1))
+let main = 2 ^^ 10 |};
+  [%expect
+    {|
+    val ^^: int -> int -> int
+    val main: int|}]
+;;
+
+let%expect_test "custom_infix_compose" =
+  pretty_printer_parse_and_infer
+    {| let ( @@ ) f g = fun x -> f (g x)
+let succ x = x + 1
+let double x = x * 2
+let main = (succ @@ double) 10 |};
+  [%expect
+    {|
+    val @@: (t3 -> t4) -> (t2 -> t3) -> t2 -> t4
+    val double: int -> int
+    val main: int
+    val succ: int -> int|}]
+;;

EML/tests/llvm.t

@@ -81,3 +81,12 @@ SPDX-License-Identifier: LGPL-3.0-or-later
 
   $ make compile_llvm additional_tests/custom_op_cat.ml
   34
+
+  $ make compile_llvm additional_tests/custom_op_right_assoc_diff.ml
+  8
+
+  $ make compile_llvm additional_tests/custom_op_via_op.ml
+  13
+
+  $ make compile_llvm additional_tests/custom_op_pipe.ml
+  11