Update AddControlled

visitor/gate.go

@@ -8,13 +8,18 @@ import (
 
 // AddControlled returns a controlled-u gate with control bit.
 // u is a (2**n x 2**n) unitary matrix and returns a (2**n x 2**n) matrix.
-func AddControlled(u matrix.Matrix, c int) matrix.Matrix {
+func AddControlled(u matrix.Matrix, c []int) matrix.Matrix {
 	d, _ := u.Dimension()
 	n := number.Log2(d)
 	g := gate.I(n)
 
+	var mask int
+	for _, bit := range c {
+		mask |= (1 << (n - 1 - bit))
+	}
+
 	for i := 0; i < d; i++ {
-		if (i>>(n-1-c))&1 == 0 {
+		if (i & mask) != mask {
 			continue
 		}
 

visitor/gate_test.go

@@ -70,7 +70,7 @@ func TestAddControlled(t *testing.T) {
 	}
 
 	for _, c := range cases {
-		got := visitor.AddControlled(c.in, c.bit)
+		got := visitor.AddControlled(c.in, []int{c.bit})
 		if !got.Equals(c.want) {
 			t.Fail()
 		}

visitor/visitor.go

@@ -316,17 +316,10 @@ func (v *Visitor) Modify(u matrix.Matrix, qargs [][]q.Qubit, modifier []parser.I
 		n := v.Visit(mod).(int64)
 		switch {
 		case mod.CTRL() != nil:
-			for _, a := range qargs[i] {
-				u = AddControlled(u, a.Index())
-			}
+			u = AddControlled(u, q.Index(qargs[i]...))
 		case mod.NEGCTRL() != nil:
-			n := v.qsim.NumberOfBit()
-			x := gate.TensorProduct(gate.X(), n, q.Index(qargs[i]...))
-
-			for _, a := range qargs[i] {
-				u = AddControlled(u, a.Index())
-			}
-
+			x := gate.TensorProduct(gate.X(), v.qsim.NumberOfBit(), q.Index(qargs[i]...))
+			u = AddControlled(u, q.Index(qargs[i]...))
 			u = matrix.Apply(x, u, x)
 		case mod.INV() != nil:
 			u = u.Dagger()