Worklist (#8)

qupsy/language.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from textwrap import indent
 
@@ -420,6 +422,9 @@ def __str__(self) -> str:
     def __repr__(self) -> str:
         return f"Pgm({self.body:!r})"
 
+    def __lt__(self, other: Pgm) -> bool:
+        return self.cost < other.cost
+
     @property
     def cost(self) -> int:
         return self.body.cost

qupsy/worklist.py

@@ -0,0 +1,30 @@
+from queue import PriorityQueue
+
+from qupsy.language import Pgm
+
+
+class Worklist:
+    def __init__(self) -> None:
+        self.current_set: PriorityQueue[tuple[int, int, Pgm]] = PriorityQueue()
+        self.overall_set: list[Pgm] = []
+
+    def put(self, *pgms: Pgm) -> None:
+        for pgm in pgms:
+            if pgm not in self.overall_set:
+                self.current_set.put((pgm.cost, pgm.depth, pgm))
+                self.overall_set.append(pgm)
+
+    def get(self) -> Pgm:
+        return self.current_set.get_nowait()[2]
+
+    def show_set(self) -> None:
+        print(self.overall_set)
+
+    def show_pq(self) -> None:
+        print(self.current_set.queue)
+
+    def notEmpty(self) -> bool:
+        return not self.current_set.empty()
+
+    def num_pgm_left(self) -> int:
+        return self.current_set.qsize()

tests/test_worklist.py

@@ -0,0 +1,30 @@
+from qupsy.language import GateCmd, H, HoleAexp, HoleCmd, HoleGate, Pgm, SeqCmd, X
+from qupsy.worklist import Worklist
+
+
+def test_worklist():
+    worklist = Worklist()
+    assert not worklist.notEmpty()
+
+
+def test_add_same_pgm():
+    worklist = Worklist()
+    pgm = Pgm(HoleCmd())
+    worklist.put(pgm, pgm)
+    assert worklist.num_pgm_left() == 1
+
+
+def test_get_pgm():
+    worklist = Worklist()
+    pgm1 = Pgm(HoleCmd())
+    pgm2 = Pgm(GateCmd(HoleGate()))
+    worklist.put(pgm1, pgm2)
+    assert worklist.get() == pgm2
+
+
+def test_add_same_cost():
+    worklist = Worklist()
+    pgm1 = Pgm(SeqCmd(GateCmd(H(HoleAexp())), GateCmd(H(HoleAexp()))))
+    pgm2 = Pgm(SeqCmd(GateCmd(X(HoleAexp())), GateCmd(X(HoleAexp()))))
+    worklist.put(pgm1, pgm2)
+    assert worklist.get() == pgm1