0 | module Syntax.StringDiagram.Braided
  1 |
  2 | import public Control.Category.Core
  3 | import public Control.Category.Functor
  4 | import public Control.Category.Monoidal
  5 | import public Control.Category.Braided
  6 | import public Data.Fin
  7 | import Data.Wrap0
  8 | import public Language.Reflection
  9 | import public Syntax.StringDiagram.Util
 10 |
 11 | %default total
 12 | %language ElabReflection
 13 |
 14 | export infix 0 -<
 15 | export prefix 0 =<
 16 |
 17 | ||| Expand shortened string diagram notation into a full expression.
 18 | |||
 19 | ||| This takes two additional parameters: `obj`, which is a wildcard
 20 | ||| expression used when passing objects as parameters, and `impl`,
 21 | ||| which is the `Braided` implementation the expansion will use.
 22 | export
 23 | stringImpl : (obj,impl : TTImp) -> TTImp -> Elab TTImp
 24 | stringImpl obj impl t = do
 25 |   MkSDiagram inp steps out <- parseDiagram t
 26 |   let Just inp' = the (Maybe _) $ sequence inp
 27 |     | Nothing => fail "Strings cannot be discarded"
 28 |   mon <- genSym "mon"
 29 |   cat <- genSym "cat"
 30 |   ts <- stringImpl' [<] (IVar EmptyFC mon) inp' out steps
 31 |   pure `(let Control.Category.Braided.MkBraided @{~(IBindVar EmptyFC mon)} {} = ~impl
 32 |              Control.Category.Monoidal.MkMonoidal @{~(IBindVar EmptyFC cat)} {} = ~(IVar EmptyFC mon)
 33 |         in ~(composeImp (IVar EmptyFC cat) ts))
 34 |   where
 35 |     listImp : Nat -> TTImp
 36 |     listImp Z = `(Prelude.Nil)
 37 |     listImp (S n) = `(Prelude.(::) ~obj ~(listImp n))
 38 |
 39 |     stringsToBack : SnocList TTImp -> Nat -> List Nat -> SnocList TTImp
 40 |     stringsToBack ts n [] = ts
 41 |     stringsToBack ts n (i :: is) =
 42 |       let iminus = pred (n `minus` i)
 43 |       in if isSucc iminus
 44 |           then stringsToBack
 45 |                 (ts :<
 46 |                   `(Control.Category.Braided.sendToBack @{~impl}
 47 |                     {xs = ~(listImp i), x = ~obj, ys = ~(listImp iminus)}))
 48 |                 n $ assert_smaller (i::is) $ map (\i' => if i' > i then pred i' else i') is
 49 |           else stringsToBack ts n is
 50 |
 51 |     stringImpl' : SnocList TTImp -> TTImp ->
 52 |                   List String -> List String -> List SDiagramStep -> Elab (SnocList TTImp)
 53 |     stringImpl' ts mon strings out (step :: steps) = do
 54 |       let Just o = the (Maybe _) $ sequence step.outputs
 55 |         | Nothing => fail "Strings cannot be discarded"
 56 |       let False = any (\name => elem name strings) o
 57 |         | True => fail "Cannot shadow string name"
 58 |       let Just is = for step.inputs $ \n =>
 59 |                       finToNat <$> findIndex (== n) strings
 60 |         | Nothing => fail "Unrecognized string name"
 61 |       let rest = filter (\n => not $ elem n step.inputs) strings
 62 |       let ts' = stringsToBack [<] (length strings) is
 63 |       stringImpl' (ts ++ ts' :<
 64 |         `(Control.Category.Monoidal.applyAssoc @{~mon}
 65 |           {xs = ~(listImp $ length rest), ys = ~(listImp $ length step.inputs),
 66 |            ys' = ~(listImp $ length o), zs = Prelude.Nil} ~(step.mor)))
 67 |         mon (rest ++ o) out steps
 68 |     stringImpl' ts _ strings out [] = do
 69 |       let Just is = for out $ \n =>
 70 |                       finToNat <$> findIndex (== n) strings
 71 |         | Nothing => fail "Unrecognized string name"
 72 |       when (length strings /= length out) $ fail "Return strings violate linearity"
 73 |       let ts' = stringsToBack [<] (length strings) is
 74 |       pure $ ts ++ ts'
 75 |
 76 |     composeImp : TTImp -> SnocList TTImp -> TTImp
 77 |     composeImp cat [<] = `(Control.Category.Core.id {a = ~obj})
 78 |     composeImp cat [<t] = t
 79 |     composeImp cat (ts :< t) =
 80 |       `(Control.Category.Core.(.) @{~cat}
 81 |         {a = ~obj, b = ~obj, c = ~obj}
 82 |         ~t ~(composeImp cat ts))
 83 |
 84 | ||| Enter string diagram notation (for `Braided`).
 85 | |||
 86 | ||| This elaboration script must be specifically invoked with
 87 | ||| `%runElab`. The category is inferred from the return type, but
 88 | ||| the tensor product must be passed as an explicit argument.
 89 | |||
 90 | ||| NOTE: This notation is intended for use with symmetric monoidal
 91 | ||| categories, and may be confusing if used in the non-symmetric
 92 | ||| braided case. If your category is non-symmetric, it may be better
 93 | ||| to use the string diagram notation for `Monoidal` and insert
 94 | ||| explicit braiding calls.
 95 | export
 96 | string : {obj : _} -> {0 cat : Hom obj} -> (0 ten : obj -> obj -> obj) -> {0 i,a,b : obj} ->
 97 |          Braided cat ten i => TTImp -> Elab (cat a b)
 98 | string {obj = obj@(Wrap0 _)} _ @{impl} t = check !(stringImpl `(W0 _) !(quote impl) t)
 99 | string {obj = _} @{impl} _ t = check !(stringImpl `(_) !(quote impl) t)
100 |