From 7e18cfe1d97d3eeee801f77d19548259311fec1c Mon Sep 17 00:00:00 2001
From: Charles Yuan <charlesyuan@google.com>
Date: Mon, 1 Jul 2024 11:47:48 -0700
Subject: [PATCH] Add explicit test for zero-size registers

---
 qualtran/_infra/registers.py      |  2 +-
 qualtran/_infra/registers_test.py | 17 +++++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/qualtran/_infra/registers.py b/qualtran/_infra/registers.py
index 2d3529f29a..8106023614 100644
--- a/qualtran/_infra/registers.py
+++ b/qualtran/_infra/registers.py
@@ -151,7 +151,7 @@ def build(cls, **registers: Union[int, sympy.Expr]) -> 'Signature':
                 will be 0-dimensional and THRU.
         """
         return cls(
-            Register(name=k, dtype=QBit() if v == 1 else QAny(v)) for k, v in registers.items() if v
+            Register(name=k, dtype=QBit() if v == 1 else QAny(v)) for k, v in registers.items()
         )
 
     @classmethod
diff --git a/qualtran/_infra/registers_test.py b/qualtran/_infra/registers_test.py
index f6dc7e5378..1a95b0d1e4 100644
--- a/qualtran/_infra/registers_test.py
+++ b/qualtran/_infra/registers_test.py
@@ -33,6 +33,17 @@ def test_register():
     assert r == r.adjoint()
 
 
+def test_zero_register():
+    r = Register("my_reg", QAny(0))
+    assert r.name == "my_reg"
+    assert r.bitsize == 0
+    assert r.shape == tuple()
+    assert r.side == Side.THRU
+    assert r.total_bits() == 0
+
+    assert r == r.adjoint()
+
+
 def test_multidim_register():
     r = Register("my_reg", QBit(), shape=(2, 3), side=Side.RIGHT)
     idxs = list(r.all_idxs())
@@ -127,6 +138,12 @@ def test_signature_build():
     sig2 = Signature.build(r1=5, r2=2)
     assert sig1 == sig2
     assert sig1.n_qubits() == 7
+
+    sig1 = Signature([Register("r1", QAny(0)), Register("r2", QAny(2))])
+    sig2 = Signature.build(r1=0, r2=2)
+    assert sig1 == sig2
+    assert sig1.n_qubits() == 2
+
     sig1 = Signature([Register("r1", QInt(7)), Register("r2", QBit())])
     sig2 = Signature.build_from_dtypes(r1=QInt(7), r2=QBit())
     assert sig1 == sig2