speed up ==, hash, <, >, <=, and >= for plain_data

[nmutil.git] / src / nmutil / lut.py

diff --git a/src/nmutil/lut.py b/src/nmutil/lut.py
index 0094bcab234f3dfa0e17ec33b389357551372682..755747ab2073dbf1a7620f9ac31e592b2bf63a44 100644 (file)
--- a/src/nmutil/lut.py
+++ b/src/nmutil/lut.py
@@ -16,7 +16,7 @@ from nmigen.hdl.ast import Array, Cat, Repl, Signal
 from nmigen.hdl.dsl import Module
 from nmigen.hdl.ir import Elaboratable
 from nmigen.cli import rtlil
-from dataclasses import dataclass
+from nmutil.plain_data import plain_data
 
 
 class BitwiseMux(Elaboratable):
@@ -83,15 +83,35 @@ class BitwiseLut(Elaboratable):
         return list(self.inputs) + [self.lut, self.output]
 
 
-@dataclass
+@plain_data()
 class _TreeMuxNode:
-    """Mux in tree for `TreeBitwiseLut`."""
+    """Mux in tree for `TreeBitwiseLut`.
+
+    Attributes:
     out: Signal
-    container: "TreeBitwiseLut"
-    parent: "_TreeMuxNode | None"
-    child0: "_TreeMuxNode | None"
-    child1: "_TreeMuxNode | None"
+    container: TreeBitwiseLut
+    parent: _TreeMuxNode | None
+    child0: _TreeMuxNode | None
+    child1: _TreeMuxNode | None
     depth: int
+    """
+    __slots__ = "out", "container", "parent", "child0", "child1", "depth"
+
+    def __init__(self, out, container, parent, child0, child1, depth):
+        """ Arguments:
+        out: Signal
+        container: TreeBitwiseLut
+        parent: _TreeMuxNode | None
+        child0: _TreeMuxNode | None
+        child1: _TreeMuxNode | None
+        depth: int
+        """
+        self.out = out
+        self.container = container
+        self.parent = parent
+        self.child0 = child0
+        self.child1 = child1
+        self.depth = depth
 
     @property
     def child_index(self):
@@ -185,3 +205,9 @@ class TreeBitwiseLut(Elaboratable):
 # useful to see what is going on:
 # python3 src/nmutil/test/test_lut.py
 # yosys <<<"read_ilang sim_test_out/__main__.TestBitwiseLut.test_tree/0.il; proc;;; show top"
+
+if __name__ == '__main__':
+    dut = BitwiseLut(2, 64)
+    vl = rtlil.convert(dut, ports=dut.ports())
+    with open("test_lut2.il", "w") as f:
+        f.write(vl)