From 4ebc3561e1773828c437261dceb32137ebdc5965 Mon Sep 17 00:00:00 2001
From: Danil Usoltsev <danilusoltsev05@gmail.com>
Date: Fri, 13 Mar 2026 12:25:36 +0300
Subject: [PATCH] add more peephole tests

---
 EML/lib/backend/ricsv/generator.ml  |  17 +-
 EML/lib/backend/ricsv/generator.mli |   1 +
 EML/lib/backend/ricsv/runner.ml     |   5 +-
 EML/lib/backend/ricsv/runner.mli    |   1 +
 EML/tests/riscv_peephole_tests.ml   | 365 ++++++++++++++++++++++++++--
 5 files changed, 359 insertions(+), 30 deletions(-)

diff --git a/EML/lib/backend/ricsv/generator.ml b/EML/lib/backend/ricsv/generator.ml
index 8dcbff5a..17c1f9ab 100644
--- a/EML/lib/backend/ricsv/generator.ml
+++ b/EML/lib/backend/ricsv/generator.ml
@@ -455,16 +455,19 @@ let bind_param_to_stack env i = function
   | _ -> fail "unsupported pattern"
 ;;
 
-let flush_instr_buffer ppf =
+let flush_instr_buffer ~enable_peephole ppf =
   let* state = get in
   let instruction_buffer = state.instr_buffer in
   let* () = put { state with instr_buffer = [] } in
-  let optimized_instructions = Peephole.optimize (List.rev instruction_buffer) in
-  let () = List.iter (fun item -> format_item ppf item) optimized_instructions in
+  let instructions = List.rev instruction_buffer in
+  let instructions =
+    if enable_peephole then Peephole.optimize instructions else instructions
+  in
+  let () = List.iter (fun item -> format_item ppf item) instructions in
   return ()
 ;;
 
-let gen_func ~enable_gc asm_name params body frame_sz ppf =
+let gen_func ~enable_gc ~enable_peephole asm_name params body frame_sz ppf =
   fprintf ppf "\n  .globl %s\n  .type %s, @function\n" asm_name asm_name;
   let args = List.length params in
   let params_reg, params_stack =
@@ -489,11 +492,11 @@ let gen_func ~enable_gc asm_name params body frame_sz ppf =
   let* () = spill_params_to_frame params_reg in
   let* () = gen_anf result_reg body in
   let* () = append (epilogue ~enable_gc ~is_main:(String.equal asm_name "main")) in
-  let* () = flush_instr_buffer ppf in
+  let* () = flush_instr_buffer ~enable_peephole ppf in
   return ()
 ;;
 
-let gen_program ~enable_gc ppf (analysis : analysis_result) =
+let gen_program ~enable_gc ~enable_peephole ppf (analysis : analysis_result) =
   fprintf ppf ".section .text";
   let base = Runtime.Primitives.runtime_primitive_arities in
   let arity_map =
@@ -519,7 +522,7 @@ let gen_program ~enable_gc ppf (analysis : analysis_result) =
       let* () =
         modify (fun state -> { state with current_func_index = function_index })
       in
-      gen_func ~enable_gc fn.asm_name fn.params fn.body frame_sz ppf)
+      gen_func ~enable_gc ~enable_peephole fn.asm_name fn.params fn.body frame_sz ppf)
   in
   match run comp init with
   | Ok ((), _) ->
diff --git a/EML/lib/backend/ricsv/generator.mli b/EML/lib/backend/ricsv/generator.mli
index 5d1bfb22..b49ab9d9 100644
--- a/EML/lib/backend/ricsv/generator.mli
+++ b/EML/lib/backend/ricsv/generator.mli
@@ -4,6 +4,7 @@
 
 val gen_program
   :  enable_gc:bool
+  -> enable_peephole:bool
   -> Format.formatter
   -> Analysis.analysis_result
   -> (unit, string) Result.t
diff --git a/EML/lib/backend/ricsv/runner.ml b/EML/lib/backend/ricsv/runner.ml
index 9efd4272..a6cb52b3 100644
--- a/EML/lib/backend/ricsv/runner.ml
+++ b/EML/lib/backend/ricsv/runner.ml
@@ -5,6 +5,7 @@
 open Middleend.Anf
 open Analysis
 
-let gen_program ?(enable_gc = false) ppf (program : anf_program) =
-  program |> analyze |> Generator.gen_program ~enable_gc ppf
+let gen_program ?(enable_gc = false) ?(enable_peephole = true) ppf (program : anf_program)
+  =
+  program |> analyze |> Generator.gen_program ~enable_gc ~enable_peephole ppf
 ;;
diff --git a/EML/lib/backend/ricsv/runner.mli b/EML/lib/backend/ricsv/runner.mli
index e5df8f59..c6ef9f26 100644
--- a/EML/lib/backend/ricsv/runner.mli
+++ b/EML/lib/backend/ricsv/runner.mli
@@ -4,6 +4,7 @@
 
 val gen_program
   :  ?enable_gc:bool
+  -> ?enable_peephole:bool
   -> Format.formatter
   -> Middleend.Anf.anf_program
   -> (unit, string) result
diff --git a/EML/tests/riscv_peephole_tests.ml b/EML/tests/riscv_peephole_tests.ml
index bf719e66..65cde2b0 100644
--- a/EML/tests/riscv_peephole_tests.ml
+++ b/EML/tests/riscv_peephole_tests.ml
@@ -3,6 +3,10 @@
 (** SPDX-License-Identifier: LGPL-3.0-or-later *)
 
 open EML_lib
+open Frontend.Parser
+open Middleend.Anf
+open Middleend.Inferencer
+open Middleend.Resolve_builtins
 open Backend.Ricsv.Architecture
 open Riscv_backend
 
@@ -13,15 +17,45 @@ let print_instrs instructions =
   print_endline (String.concat "\n" rendered)
 ;;
 
-let run_peephole input = input |> Backend.Ricsv.Peephole.optimize |> print_instrs
+let compile_riscv ?(enable_peephole = true) src =
+  match parse src with
+  | Error e -> "Parse error: " ^ e
+  | Ok ast ->
+    let scope = TypeEnv.keys TypeEnv.initial_env in
+    let ast' = resolve_program ast scope in
+    (match anf_program ast' with
+     | Error e -> "ANF error: " ^ e
+     | Ok anf ->
+       let buf = Buffer.create 1024 in
+       let ppf = Format.formatter_of_buffer buf in
+       (match Backend.Ricsv.Runner.gen_program ~enable_peephole ppf anf with
+        | Ok () ->
+          Format.pp_print_flush ppf ();
+          Buffer.contents buf
+        | Error e -> "Codegen error: " ^ e))
+;;
 
-(** Печатает код без пипхолов и с пипхолами, чтобы видеть разницу. *)
-let run_peephole_show_diff input =
+let show_diff ~input ~output value =
   print_endline "=== Without peepholes ===";
-  print_instrs input;
+  input value;
   print_endline "";
   print_endline "=== With peepholes ===";
-  run_peephole input
+  output value
+;;
+
+let show_with_without_peephole src =
+  show_diff
+    ~input:(fun source -> print_endline (compile_riscv ~enable_peephole:false source))
+    ~output:(fun source -> print_endline (compile_riscv ~enable_peephole:true source))
+    src
+;;
+
+let show_peephole_diff instrs =
+  show_diff
+    ~input:print_instrs
+    ~output:(fun instructions ->
+      instructions |> Backend.Ricsv.Peephole.optimize |> print_instrs)
+    instrs
 ;;
 
 let%expect_test "optimizes repeated stack load pattern from task description" =
@@ -36,7 +70,7 @@ let%expect_test "optimizes repeated stack load pattern from task description" =
     ; Sd (T 1, (SP, 64))
     ]
   in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -60,7 +94,7 @@ let%expect_test "removes redundant load and forwards store to load" =
   let input =
     [ Ld (T 0, (SP, 64)); Ld (T 1, (SP, 64)); Sd (T 1, (SP, 64)); Ld (A 0, (SP, 64)) ]
   in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -81,7 +115,7 @@ let%expect_test "folds addi chain and removes dead overwrite" =
   let input =
     [ Addi (SP, SP, -16); Addi (SP, SP, 8); Li (T 0, 1); Li (T 0, 2); Addi (T 1, T 1, 0) ]
   in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -99,7 +133,7 @@ li t0, 2
 
 let%expect_test "drops jump to the immediately following label" =
   let input = [ J "l1"; Label "l1"; Mv (A 0, A 0); Ret ] in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -115,7 +149,7 @@ ret
 
 let%expect_test "collapses double copy before binary op" =
   let input = [ Mv (T 0, A 0); Mv (T 1, A 0); Add (A 0, T 0, T 1) ] in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -130,7 +164,7 @@ add a0, a0, a0
 
 let%expect_test "propagates single mv into following consumer" =
   let input = [ Mv (T 0, A 0); Li (T 1, 1); Slt (A 0, T 0, T 1) ] in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -146,7 +180,7 @@ slt a0, a0, t1
 
 let%expect_test "rewrites li plus add into addi" =
   let input = [ Li (T 1, 1); Add (A 0, T 0, T 1) ] in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -160,7 +194,7 @@ addi a0, t0, 1
 
 let%expect_test "folds li plus add when destination is constant register" =
   let input = [ Li (T 0, 1); Add (T 0, T 1, T 0) ] in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -174,7 +208,7 @@ addi t0, t1, 1
 
 let%expect_test "rewrites mul by power of two into slli" =
   let input = [ Li (T 0, 4); Mul (A 0, T 1, T 0) ] in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -190,7 +224,7 @@ let%expect_test "keeps load cache barriers on call" =
   let input =
     [ Ld (T 0, (SP, 64)); Call "foo"; Ld (T 1, (SP, 64)); Add (A 0, T 0, T 1) ]
   in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -209,7 +243,7 @@ add a0, t0, t1
 
 let%expect_test "forwards store to following load on same slot" =
   let input = [ Sd (A 0, (fp, -16)); Ld (T 0, (fp, -16)); Add (A 0, T 0, A 1) ] in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -225,7 +259,7 @@ add a0, a0, a1
 
 let%expect_test "folds constant beq into jump" =
   let input = [ Li (T 0, 1); Li (T 1, 1); Beq (T 0, T 1, "else_1") ] in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -242,7 +276,7 @@ j else_1
 
 let%expect_test "removes dead store before ret in same block" =
   let input = [ Sd (A 0, (fp, -8)); Add (A 0, A 0, A 1); Ret ] in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -258,7 +292,7 @@ ret
 
 let%expect_test "keeps store before call barrier" =
   let input = [ Sd (A 0, (fp, -8)); Call "foo"; Ret ] in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -275,7 +309,7 @@ ret
 
 let%expect_test "removes store that restores unchanged loaded slot value" =
   let input = [ Ld (T 1, (sp, 64)); Add (T 0, T 1, T 0); Sd (T 1, (sp, 64)) ] in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -291,7 +325,7 @@ add t0, t1, t0
 
 let%expect_test "drops overwritten store before next store to same slot" =
   let input = [ Sd (A 0, (fp, -8)); Add (A 0, A 0, A 1); Sd (T 0, (fp, -8)); Ret ] in
-  run_peephole_show_diff input;
+  show_peephole_diff input;
   [%expect
     {|
 === Without peepholes ===
@@ -305,3 +339,292 @@ add a0, a0, a1
 ret
 |}]
 ;;
+
+let%expect_test "shows code with and without peephole 1" =
+  let src =
+    {|
+      let f x =
+        let y = x < 0 in
+        y + 1
+      let main = f 1
+    |}
+  in
+  show_with_without_peephole src;
+  [%expect
+    {|
+    === Without peepholes ===
+    .section .text
+      .globl f
+      .type f, @function
+    f:
+      addi sp, sp, -24
+      sd ra, 16(sp)
+      sd fp, 8(sp)
+      addi fp, sp, 8
+      sd a0, -8(fp)
+      ld t0, -8(fp)
+      li t1, 1
+      slt a0, t0, t1
+      add a0, a0, a0
+      addi a0, a0, 1
+      sd a0, -16(fp)
+      ld t0, -16(fp)
+      li t1, 3
+      add a0, t0, t1
+      addi a0, a0, -1
+      addi sp, fp, 16
+      ld ra, 8(fp)
+      ld fp, 0(fp)
+      ret
+
+      .globl main
+      .type main, @function
+    main:
+      addi sp, sp, -200
+      sd ra, 192(sp)
+      sd fp, 184(sp)
+      addi fp, sp, 184
+      li a0, 3
+      call f
+      addi sp, fp, 16
+      ld ra, 8(fp)
+      ld fp, 0(fp)
+      li a0, 0
+      ret
+
+
+    === With peepholes ===
+    .section .text
+      .globl f
+      .type f, @function
+    f:
+      addi sp, sp, -24
+      sd ra, 16(sp)
+      sd fp, 8(sp)
+      addi fp, sp, 8
+      sd a0, -8(fp)
+      li t1, 1
+      slt a0, a0, t1
+      add a0, a0, a0
+      addi a0, a0, 1
+      sd a0, -16(fp)
+      addi a0, a0, 2
+      addi sp, fp, 16
+      ld ra, 8(fp)
+      ld fp, 0(fp)
+      ret
+
+      .globl main
+      .type main, @function
+    main:
+      addi sp, sp, -200
+      sd ra, 192(sp)
+      sd fp, 184(sp)
+      addi fp, sp, 184
+      li a0, 3
+      call f
+      addi sp, fp, 16
+      ld ra, 8(fp)
+      ld fp, 0(fp)
+      li a0, 0
+      ret |}]
+;;
+
+let%expect_test "shows code with and without peephole 2" =
+  let src =
+    {|
+      let f y =
+        let x = 1 + y in
+        let z = 2 * y in
+        x + z
+      let main = f 10
+    |}
+  in
+  show_with_without_peephole src;
+  [%expect
+    {|
+    === Without peepholes ===
+    .section .text
+      .globl f
+      .type f, @function
+    f:
+      addi sp, sp, -32
+      sd ra, 24(sp)
+      sd fp, 16(sp)
+      addi fp, sp, 16
+      sd a0, -8(fp)
+      li t0, 3
+      ld t1, -8(fp)
+      add a0, t0, t1
+      addi a0, a0, -1
+      sd a0, -16(fp)
+      li t0, 5
+      ld t1, -8(fp)
+      srli t0, t0, 1
+      addi t1, t1, -1
+      mul a0, t0, t1
+      addi a0, a0, 1
+      sd a0, -24(fp)
+      ld t0, -16(fp)
+      ld t1, -24(fp)
+      add a0, t0, t1
+      addi a0, a0, -1
+      addi sp, fp, 16
+      ld ra, 8(fp)
+      ld fp, 0(fp)
+      ret
+
+      .globl main
+      .type main, @function
+    main:
+      addi sp, sp, -200
+      sd ra, 192(sp)
+      sd fp, 184(sp)
+      addi fp, sp, 184
+      li a0, 21
+      call f
+      addi sp, fp, 16
+      ld ra, 8(fp)
+      ld fp, 0(fp)
+      li a0, 0
+      ret
+
+
+    === With peepholes ===
+    .section .text
+      .globl f
+      .type f, @function
+    f:
+      addi sp, sp, -32
+      sd ra, 24(sp)
+      sd fp, 16(sp)
+      addi fp, sp, 16
+      sd a0, -8(fp)
+      li t0, 3
+      mv t1, a0
+      add a0, t0, t1
+      addi a0, a0, -1
+      sd a0, -16(fp)
+      li t0, 5
+      srli t0, t0, 1
+      addi t1, t1, -1
+      mul a0, t0, t1
+      addi a0, a0, 1
+      sd a0, -24(fp)
+      ld t0, -16(fp)
+      add a0, t0, a0
+      addi a0, a0, -1
+      addi sp, fp, 16
+      ld ra, 8(fp)
+      ld fp, 0(fp)
+      ret
+
+      .globl main
+      .type main, @function
+    main:
+      addi sp, sp, -200
+      sd ra, 192(sp)
+      sd fp, 184(sp)
+      addi fp, sp, 184
+      li a0, 21
+      call f
+      addi sp, fp, 16
+      ld ra, 8(fp)
+      ld fp, 0(fp)
+      li a0, 0
+      ret |}]
+;;
+
+let%expect_test "shows code with and without peephole 3" =
+  let src =
+    {|
+      let g x =
+        let a = x + 1 in
+        let b = a + 1 in
+        b + 1
+      let main = g 1
+    |}
+  in
+  show_with_without_peephole src;
+  [%expect
+    {|
+    === Without peepholes ===
+    .section .text
+      .globl g
+      .type g, @function
+    g:
+      addi sp, sp, -32
+      sd ra, 24(sp)
+      sd fp, 16(sp)
+      addi fp, sp, 16
+      sd a0, -8(fp)
+      ld t0, -8(fp)
+      li t1, 3
+      add a0, t0, t1
+      addi a0, a0, -1
+      sd a0, -16(fp)
+      ld t0, -16(fp)
+      li t1, 3
+      add a0, t0, t1
+      addi a0, a0, -1
+      sd a0, -24(fp)
+      ld t0, -24(fp)
+      li t1, 3
+      add a0, t0, t1
+      addi a0, a0, -1
+      addi sp, fp, 16
+      ld ra, 8(fp)
+      ld fp, 0(fp)
+      ret
+
+      .globl main
+      .type main, @function
+    main:
+      addi sp, sp, -200
+      sd ra, 192(sp)
+      sd fp, 184(sp)
+      addi fp, sp, 184
+      li a0, 3
+      call g
+      addi sp, fp, 16
+      ld ra, 8(fp)
+      ld fp, 0(fp)
+      li a0, 0
+      ret
+
+
+    === With peepholes ===
+    .section .text
+      .globl g
+      .type g, @function
+    g:
+      addi sp, sp, -32
+      sd ra, 24(sp)
+      sd fp, 16(sp)
+      addi fp, sp, 16
+      sd a0, -8(fp)
+      addi a0, a0, 2
+      sd a0, -16(fp)
+      addi a0, a0, 2
+      sd a0, -24(fp)
+      addi a0, a0, 2
+      addi sp, fp, 16
+      ld ra, 8(fp)
+      ld fp, 0(fp)
+      ret
+
+      .globl main
+      .type main, @function
+    main:
+      addi sp, sp, -200
+      sd ra, 192(sp)
+      sd fp, 184(sp)
+      addi fp, sp, 184
+      li a0, 3
+      call g
+      addi sp, fp, 16
+      ld ra, 8(fp)
+      ld fp, 0(fp)
+      li a0, 0
+      ret |}]
+;;