From a92f01db206b761ebf7f26c90eb5ae54266b8cca Mon Sep 17 00:00:00 2001
From: Michel Weber <michel.weber@inf.ethz.ch>
Date: Wed, 11 Jan 2023 20:03:07 +0100
Subject: [PATCH] tests: Add test for xdsl_opt passes

---
 tests/xdsl_opt/constant_program.xdsl |  3 +++
 tests/xdsl_opt/test_xdsl_opt.py      | 26 ++++++++++++++++++++++++++
 xdsl/xdsl_opt_main.py                | 10 +++++++---
 3 files changed, 36 insertions(+), 3 deletions(-)
 create mode 100644 tests/xdsl_opt/constant_program.xdsl

diff --git a/tests/xdsl_opt/constant_program.xdsl b/tests/xdsl_opt/constant_program.xdsl
new file mode 100644
index 0000000000..4af2312d54
--- /dev/null
+++ b/tests/xdsl_opt/constant_program.xdsl
@@ -0,0 +1,3 @@
+builtin.module() {
+  %0 : !i32 = arith.constant() ["value" = 5 : !i32]
+}
diff --git a/tests/xdsl_opt/test_xdsl_opt.py b/tests/xdsl_opt/test_xdsl_opt.py
index 90a7ef35b5..55dcf6896b 100644
--- a/tests/xdsl_opt/test_xdsl_opt.py
+++ b/tests/xdsl_opt/test_xdsl_opt.py
@@ -3,6 +3,9 @@
 from contextlib import redirect_stdout
 from io import StringIO
 import os
+from xdsl.ir import MLContext
+from xdsl.dialects.builtin import ModuleOp
+from xdsl.rewriter import Rewriter
 
 
 def test_opt():
@@ -72,3 +75,26 @@ def test_print_to_file():
     expected = open(filename_out, 'r').read()
 
     assert inp.strip() == expected.strip()
+
+
+def test_operation_deletion():
+    filename_in = 'tests/xdsl_opt/constant_program.xdsl'
+    filename_out = 'tests/xdsl_opt/empty_program.xdsl'
+
+    class xDSLOptMainPass(xDSLOptMain):
+
+        def register_all_passes(self):
+
+            def remove_constant(ctx: MLContext, module: ModuleOp):
+                module.ops[0].detach()
+
+            self.available_passes['remove-constant'] = remove_constant
+
+    opt = xDSLOptMainPass(args=[filename_in, '-p', 'remove-constant'])
+
+    f = StringIO("")
+    with redirect_stdout(f):
+        opt.run()
+    expected = open(filename_out, 'r').read()
+
+    assert f.getvalue().strip() == expected.strip()
diff --git a/xdsl/xdsl_opt_main.py b/xdsl/xdsl_opt_main.py
index 62701eb810..28e49dfdec 100644
--- a/xdsl/xdsl_opt_main.py
+++ b/xdsl/xdsl_opt_main.py
@@ -33,18 +33,18 @@ class xDSLOptMain:
     attributes.
     """
 
-    available_frontends: Dict[str, Callable[[IOBase], ModuleOp]] = {}
+    available_frontends: Dict[str, Callable[[IOBase], ModuleOp]]
     """
     A mapping from file extension to a frontend that can handle this
     file type.
     """
 
-    available_passes: Dict[str, Callable[[MLContext, ModuleOp], None]] = {}
+    available_passes: Dict[str, Callable[[MLContext, ModuleOp], None]]
     """
     A mapping from pass names to functions that apply the pass to a  ModuleOp.
     """
 
-    available_targets: Dict[str, Callable[[ModuleOp, IOBase], None]] = {}
+    available_targets: Dict[str, Callable[[ModuleOp, IOBase], None]]
     """
     A mapping from target names to functions that serialize a ModuleOp into a
     stream.
@@ -56,6 +56,10 @@ class xDSLOptMain:
     def __init__(self,
                  description: str = 'xDSL modular optimizer driver',
                  args=None):
+        self.available_frontends = {}
+        self.available_passes = {}
+        self.available_targets = {}
+
         self.ctx = MLContext()
         self.register_all_dialects()
         self.register_all_frontends()