From c6c6f83b27db9a8659af499017ec8fe4b71b00f0 Mon Sep 17 00:00:00 2001
From: Craig Roy <croyzor@users.noreply.github.com>
Date: Thu, 24 Oct 2024 09:57:48 +0100
Subject: [PATCH] docs: Add example of magic state distillation (#40)

Co-authored-by: Alan Lawrence <alan.lawrence@quantinuum.com>
---
 README.md                                   |  1 +
 brat/examples/lib/functional.brat           | 12 ++++
 brat/examples/lib/math.brat                 |  6 ++
 brat/examples/magic-state-distillation.brat | 71 +++++++++++++++++++++
 brat/examples/rus.brat                      | 14 ++--
 brat/test/Test/Compile/Hugr.hs              |  3 +-
 6 files changed, 98 insertions(+), 9 deletions(-)
 create mode 100644 brat/examples/lib/functional.brat
 create mode 100644 brat/examples/lib/math.brat
 create mode 100644 brat/examples/magic-state-distillation.brat

diff --git a/README.md b/README.md
index 7acb4e47..ccf21221 100644
--- a/README.md
+++ b/README.md
@@ -34,6 +34,7 @@ The [`brat/examples`](brat/examples) directory contains some examples of BRAT pr
 For example:
 - Quantum teleportation: [teleportation.brat](brat/examples/teleportation.brat)
 - Quantum Fourier transform: [teleportation.brat](brat/examples/teleportation.brat)
+- Magic State Distillation: [magic-state-distillation.brat](brat/examples/magic-state-distillation.brat)
 - Simple Repeat-Until-Success: [rus.brat](brat/examples/rus.brat)
 - Ising Hamiltonian: [ising.brat](brat/examples/ising.brat)
 
diff --git a/brat/examples/lib/functional.brat b/brat/examples/lib/functional.brat
new file mode 100644
index 00000000..f8ccccab
--- /dev/null
+++ b/brat/examples/lib/functional.brat
@@ -0,0 +1,12 @@
+-- TODO: Fill this with holes once we can guess them
+map(X :: $,  Y :: $, n :: #, f :: { X -o Y }) -> { Vec(X, n) -o Vec(Y, n) }
+map(_, _, _, _) = { [] => [] }
+map(X, Y, succ(n), f) = { cons(x,xs) => cons(f(x), map(X, Y, n, f)(xs)) }
+
+fold(X :: $
+    ,{ X, X -o X }
+    ,{ -o X }
+    ,n :: #
+    ) -> { Vec(X, n) -o X }
+fold(X, f, x, 0) = { [] => x() }
+fold(X, f, x, succ(n)) = { cons(y, ys) => f(y, fold(X, f, x, n)(ys)) }
diff --git a/brat/examples/lib/math.brat b/brat/examples/lib/math.brat
new file mode 100644
index 00000000..e87f3c4b
--- /dev/null
+++ b/brat/examples/lib/math.brat
@@ -0,0 +1,6 @@
+pi :: Float
+pi = 3.1415926
+
+-- Things that we need to find a way to support
+ext "TODO" arccos :: { Float -> Float }
+ext "TODO_sqrt" sqrt :: { Float -> Float }
diff --git a/brat/examples/magic-state-distillation.brat b/brat/examples/magic-state-distillation.brat
new file mode 100644
index 00000000..069551c0
--- /dev/null
+++ b/brat/examples/magic-state-distillation.brat
@@ -0,0 +1,71 @@
+open import lib.functional
+open import lib.kernel
+open import lib.math
+
+or(Bit, Bit) -o Bit
+or(true, _) = true
+or(_, y) = y
+
+any(n :: #) -> { Vec(Bit, n) -o Bit }
+any(n) = fold(Bit, or, { => false }, n)
+
+-- Decoder for the [[5,3,1]] "perfect" error code
+-- TODO: Add vectorisation and fanin/out
+decoder(Vec(Qubit, 5)) -o Vec(Qubit, 5)
+decoder = {
+  [/\];
+  CZ, CZ, |;
+  |, CZ, CZ;
+  swap, |, swap;
+  |, swap, ..;
+  |, |, CZ, |;
+  |, swap, ..;
+  swap, |, swap;
+  H, H, H, H, H;
+  [\/]
+}
+
+phi :: Float
+phi = arccos(1.0 / sqrt(3.0))
+
+prepare_noisy_1(Money) -o Qubit
+prepare_noisy_1 = {
+  fresh;
+  Ry(phi);
+  Rz(pi / 4.0)
+}
+
+measure_all(n :: #) -> { Vec(Qubit, n) -o Vec(Bit, n), Vec(Money, n) }
+measure_all(0) = { [] => [], [] }
+measure_all(succ(n)) = { q ,- qs =>
+  let m, b = measure(q) in
+    let bs, ms = measure_all(n)(qs) in
+      b ,- bs, m ,- ms
+}
+
+measure_syndrome(n :: #)
+             -> { Vec(Qubit, n)
+               -o Bit, Vec(Money, n)
+                }
+measure_syndrome(n) = {
+  measure_all(n);
+  any(n), |
+}
+
+prepare_noisy(n :: #) -> { Vec(Money, n) -o Vec(Qubit, n) }
+prepare_noisy(n) = map(Money, Qubit, n, prepare_noisy_1)
+
+redistill_if_bad(Qubit, Bit, Vec(Money, 4))
+              -o Qubit, Vec(Money, 4)
+redistill_if_bad(q, false, ms) = q, ms
+redistill_if_bad(q, true, ms) = let m, _ = measure(q) in
+                                  distill(cons(m, ms))
+
+distill(Vec(Money, 5)) -o Qubit, Vec(Money, 4)
+distill = {
+  prepare_noisy(5);
+  decoder;
+  uncons(4);
+  |, measure_syndrome(4);
+  redistill_if_bad
+}
diff --git a/brat/examples/rus.brat b/brat/examples/rus.brat
index 909a3a2c..de20867e 100644
--- a/brat/examples/rus.brat
+++ b/brat/examples/rus.brat
@@ -1,8 +1,8 @@
 open import lib.kernel
 
-
-rus(Money, Qubit) -o Money, Qubit
-rus = {
+rus(timeout :: Nat) -> { Money, Qubit -o Money, Qubit }
+rus(0) = { .. }
+rus(succ(n)) = {
   fresh, ..;
   (H;T), ..;
       CX;
@@ -10,9 +10,7 @@ rus = {
       CX;
   (T;H), ..;
   measure, ..;
-  rus'
+  ( m, false, q => m, q
+  | m, true, q => rus(n)(m, q)
+  )
 }
-
-rus'(Money, Bit, Qubit) -o Money, Qubit
-rus'(m, false, q) = m, q      -- success
-rus'(m, true, q) = rus(m, q)  -- try again
diff --git a/brat/test/Test/Compile/Hugr.hs b/brat/test/Test/Compile/Hugr.hs
index 2177c92c..5fbb79de 100644
--- a/brat/test/Test/Compile/Hugr.hs
+++ b/brat/test/Test/Compile/Hugr.hs
@@ -61,13 +61,14 @@ nonCompilingExamples = (expectedCheckingFails ++ expectedParsingFails ++
   -- signatures causes `kindCheck` to call `abstract`, creating "Selector"
   -- nodes, which we don't attempt to compile because we want to get rid of them
   ,"vec-pats"
-  -- Victims of #389
+  -- Victims of #13
   ,"arith"
   ,"bell"
   ,"cqcconf"
   ,"imports"
   ,"ising"
   ,"klet"
+  ,"magic-state-distillation" -- also makes selectors
   ,"rus"
   ,"teleportation"
   ,"vlup_covering"