WIP: Hugr compilation uses nonlocals and does not undo

brat/Brat/Compile/Hugr.hs

@@ -25,12 +25,13 @@
 import Control.Monad.Freer
 import Data.Hugr
 import Hasochism
+import Util
 
 import Control.Monad (unless)
 import Data.Aeson
 import Data.Bifunctor (first, second)
 import qualified Data.ByteString.Lazy as BS
-import Data.Foldable (traverse_, for_)
+import Data.Foldable (for_)
 import Data.Functor ((<&>), ($>))
 import Data.List (partition, sort, sortBy)
 import qualified Data.Map as M
@@ -131,7 +132,7 @@
   whoAmI = gets (fst . nameSupply)
 
 runCheckingInCompile :: Free CheckingSig t -> Compile t
 runCheckingInCompile (Ret t) = pure t
 runCheckingInCompile (Req (ELup e) k) = do
   emap <- gets (valueMap . store)
   runCheckingInCompile (k (M.lookup e emap))
@@ -285,13 +286,13 @@
   let TestMatchData my matchSeq = matchData
   matchSeq <- compileGraphTypes (fmap (binderToValue my) matchSeq)
 
-  let portTbl = zip (fst <$> matchInputs matchSeq) ins
+  let portTbl = fromJust (zipSameLength (fst <$> matchInputs matchSeq) ins)
   testResult <- compileMatchSequence parent portTbl matchSeq
 
   -- Feed the test result into a conditional
-  makeConditional ("clause of " ++ show rhs) parent testResult [] [("didntMatch", didntMatch outTys)
-                                                                  ,("didMatch", didMatch outTys)
-                                                                  ]
+  makeConditional ("clause of " ++ show rhs) parent testResult [("didntMatch", didntMatch outTys)
+                                                               ,("didMatch", didMatch outTys)
+                                                               ]
  where
   didntMatch :: [HugrType] -> NodeId -> [TypedPort] -> Compile [TypedPort]
   didntMatch outTys parent ins = case nonEmpty clauses of
@@ -611,18 +612,15 @@
 -- Return a sum whose first component is the types we started with in the order
 -- specified by the portTable argument.
 --
--- In the happy path we return wires in the order of `matchOutputs`
+-- In the happy path we return wires in the order of `rhsInputs`
 -- otherwise, the order is the same as how they came in via the portTable
 compileMatchSequence :: NodeId -- The parent node
                      -> [(Src  -- Things we've matched or passed through, coming from an Input node
                          ,TypedPort)] -- This portTable and `matchInputs` should be in the same order
                      -> MatchSequence HugrType
                      -> Compile TypedPort
 compileMatchSequence parent portTable (MatchSequence {..}) = do
-  unless
-    ((second snd <$> portTable) == matchInputs)
-    (error "compileMatchSequence assert failed")
-  let sumTy = SoR [snd <$> matchInputs, snd <$> matchOutputs]
+  let sumTy = SoR [snd <$> matchInputs, snd <$> rhsInputs]
   case matchTests of
     (src, primTest):tests -> do
       -- Pick the port corresponding to the source we want to test
@@ -633,17 +631,20 @@
       -- other inputs
       testResult <- compilePrimTest parent typedPort primTest
       let testIx = length left
-      let remainingMatchTests = MatchSequence (primTestOuts primTest ++ (second snd <$> others)) tests matchOutputs
-      ports <- makeConditional ("matching " ++ show (src, primTest)) parent testResult (snd <$> others)
-               [("didNotMatch", didNotMatchCase testIx sumTy)
+      let remainingMatchTests = MatchSequence (primTestOuts primTest ++ (second snd <$> others)) tests rhsInputs
+      ports <- makeConditional ("matching " ++ show (src, primTest)) parent testResult
+               [("didNotMatch", \parent _ -> makeRowTag "DidNotMatch" parent 0 sumTy (reorderPortTbl portTable (fst <$> matchInputs)))
                ,("didMatch",    didMatchCase testIx (primTest, snd typedPort) remainingMatchTests sumTy)]
       case ports of
         (port:_) -> pure port
         _ -> error $ "Expected at least one output port from makeConditional: got\n  " ++ show ports
 
     [] -> do
-      -- Reorder into `matchOutputs` order
-      let ins = reorderPortTbl portTable (fst <$> matchOutputs)
+      -- Reorder into `rhsInputs` order
+      rhsInputsRefinedFromUnification <- for rhsInputs $ \input@(src, hugrTy) -> compileWithInputs parent (endName (toEnd src)) >>= \case
+        Nothing -> error $ "Failed to compile rhsInput: " ++ show input
+        Just nodeId -> pure (src, (Port nodeId 0, hugrTy))
+      let ins = reorderPortTbl (portTable <> rhsInputsRefinedFromUnification) (fst <$> rhsInputs)
       -- Need to pack inputs into a tuple before feeding them into a tag node
       ports <- makeRowTag "Success" parent 1 sumTy ins
       case ports of
@@ -662,10 +663,12 @@
                -> [TypedPort]
                -> Compile [TypedPort]
   didMatchCase ix (prevTest, oldTy) ms@(MatchSequence{..}) sumTy parent ins = do
-    -- Remember which port a src corresponds to
-    let portTable = zip (fst <$> matchInputs) ins
-    didAllTestsSucceed <- compileMatchSequence parent portTable ms
-    makeConditional ("all matched (" ++ show ix ++ ")") parent didAllTestsSucceed []
+    -- We want to add the stuff generated by tests to the port table
+    let extraInputs = case prevTest of
+          PrimCtorTest _ _ _ newSrcs -> zip (fst <$> newSrcs) ins
+          PrimLitTest _ -> []
+    didAllTestsSucceed <- compileMatchSequence parent (portTable <> extraInputs) ms
+    makeConditional ("all matched (" ++ show ix ++ ")") parent didAllTestsSucceed
       [("Undo", undo)
       ,("AllMatched", allMatched)
       ]
@@ -676,29 +679,16 @@
          -> Compile [TypedPort]
     undo parent ins = do
       -- Test results, and the rest of the inputs
-      let (refined, other) = splitAt (length (primTestOuts prevTest)) ins
+      let (refined, others) = splitAt (length (primTestOuts prevTest)) ins
       undoPort <- undoPrimTest parent refined oldTy prevTest
       -- Put it back in the right place
-      let (as, bs) = splitAt ix other
+      let (as, bs) = splitAt ix others
       let ins = as ++ undoPort : bs
       makeRowTag "Fail_Undo" parent 0 sumTy ins
 
     allMatched :: NodeId -> [TypedPort] -> Compile [TypedPort]
     allMatched parent = makeRowTag "AllMatched" parent 1 sumTy
 
-  didNotMatchCase :: Int -- The index at which to put the thing we inspected in outputs
-                  -> SumOfRows
-                  -> NodeId
-                  -> [TypedPort]
-                  -> Compile [TypedPort]
-  didNotMatchCase _ _ _ [] = error "No scrutinee input in didNotMatchCase"
-  didNotMatchCase ix sumTy parent (scrutinee:ins) = do
-    let (as, bs) = splitAt ix ins
-    -- We need to wire inputs to a `Tag0`, but bringing the tested src back to
-    -- the original position
-    let ins = as ++ scrutinee:bs
-    makeRowTag "DidNotMatch" parent 0 sumTy ins
-
 makeRowTag :: String -> NodeId -> Int -> SumOfRows -> [TypedPort] -> Compile [TypedPort]
 makeRowTag hint parent tag sor@(SoR sumRows) ins =
   if sumRows !! tag == (snd <$> ins)
@@ -724,20 +714,19 @@
 makeConditional :: String    -- Label
                 -> NodeId    -- Parent node id
                 -> TypedPort -- The discriminator
-                -> [TypedPort] -- Other inputs
                 -> [(String, NodeId -> [TypedPort] -> Compile [TypedPort])] -- Must be ordered
                 -> Compile [TypedPort]
-makeConditional lbl parent discrim otherInputs cases = do
+makeConditional lbl _ discrim cases | track ("makeConditional(" ++ show lbl ++ ")\n  " ++ unlines (fst <$> cases) ++ "\n  " ++ show (snd discrim)) False = undefined
+makeConditional lbl parent discrim cases = do
   condId <- freshNode "Conditional"
   let rows = getSumVariants (snd discrim)
-  outTyss <- for (zip (zip [0..] cases) rows) (\((ix, (name, f)), row) -> makeCase condId name ix (row ++ (snd <$> otherInputs)) f)
+  outTyss <- for (fromJust $ zipSameLength (zip [0..] cases) rows) (\((ix, (name, f)), row) -> makeCase condId name ix row f)
   unless
     (allRowsEqual outTyss)
     (error "Conditional output types didn't match")
-  let condOp = OpConditional (Conditional parent rows (snd <$> otherInputs) (head outTyss) [("label", lbl)])
+  let condOp = OpConditional (Conditional parent rows [] (head outTyss) [("label", lbl)])
   addOp condOp condId
   addEdge (fst discrim, Port condId 0)
-  traverse_ addEdge (zip (fst <$> otherInputs) (Port condId <$> [1..]))
   pure $ zip (Port condId <$> [0..]) (head outTyss)
  where
   makeCase :: NodeId -> String -> Int -> [HugrType] -> (NodeId -> [TypedPort] -> Compile [TypedPort]) -> Compile [HugrType]
@@ -761,6 +750,8 @@
                 -> PrimTest HugrType -- The test to run
                 -> Compile TypedPort
 compilePrimTest parent (port, ty) (PrimCtorTest c tycon unpackingNode outputs) = do
+  -- PrimCtorTest returns the `outputs` specified in the happy case, else the
+  -- thing that was originally being tested, unchanged.
   let sumOut = HTSum (SG (GeneralSum [[ty], snd <$> outputs]))
   let sig = FunctionType [ty] [sumOut] ["BRAT"]
   testId <- addNode ("PrimCtorTest " ++ show c)
@@ -778,7 +769,8 @@
   constId <- addNode "LitConst" (OpConst (ConstOp parent (valFromSimple tm)))
   loadPort <- head <$> addNodeWithInputs "LitLoad" (OpLoadConstant (LoadConstantOp parent ty))
                        [(Port constId 0, ty)] [ty]
-  -- Connect to a test node
+  -- PrimLitTest returns an empty sum in the happy case (no extra info), else
+  -- the thing that was originally being tested, unchanged.
   let sumOut = HTSum (SG (GeneralSum [[ty], []]))
   let sig = FunctionType [ty, ty] [sumOut] ["BRAT"]
   head <$> addNodeWithInputs ("PrimLitTest " ++ show tm)
@@ -789,6 +781,12 @@
 constructorOp :: NodeId -> QualName -> QualName -> FunctionType -> HugrOp NodeId
 constructorOp parent tycon c sig = OpCustom (CustomOp parent "BRAT" ("Ctor::" ++ show tycon ++ "::" ++ show c) sig [])
 
+-- undoPrimTest, reconstructs the scrutinee of a test after the test has been run.
+-- For literal tests, the literal is recorded in the test, so we can just put it
+-- in a const node.
+-- For constructor tests, we assume that the outputs of the test are the
+-- arguments of the constructor, so we can apply them to the constructor in the
+-- same order to reproduce the scrutinee.
 undoPrimTest :: NodeId
              -> [TypedPort] -- The inputs we have to put back together
              -> HugrType -- The type of the thing we're making

brat/Brat/Graph.hs

@@ -59,19 +59,22 @@ deriving instance Show (NodeType a)
 -- tag 0 with the function's inputs returned as they were
 -- tag 1 with the environment of pattern variables from a successful
 data TestMatchData (m :: Mode) where
-  TestMatchData :: Show (BinderType m) => Modey m -> MatchSequence (BinderType m) -> TestMatchData m
+  TestMatchData :: Show (BinderType m)
+                => Modey m
+                -> MatchSequence (BinderType m)
+                -> TestMatchData m
 
 deriving instance Show (TestMatchData a)
 
 -- A collections of tests to determine if a clause matches.
 -- Invariants:
 --    1. Each src in `matchTests` has been mentioned earlier (either in `matchInputs`
 --       or in the srcs outputted by a previous `PrimCtorTest`
---    2. The same goes for the sources in `matchOutputs`
+--    2. The same goes for the sources in `rhsInputs`
 data MatchSequence ty = MatchSequence
   { matchInputs :: [(Src, ty)]
   , matchTests :: [(Src, PrimTest ty)]
-  , matchOutputs ::[(Src, ty)]
+  , rhsInputs ::[(Src, ty)]
   } deriving (Foldable, Functor, Traversable)
 deriving instance Show ty => Show (MatchSequence ty)
 

brat/test/Test/Compile/Hugr.hs

@@ -23,10 +23,6 @@ invalidExamples = map ((++ ".brat") . ("examples" </>))
   ,"app"
   ,"dollar_kind"
   ,"portpulling"
-  ,"eatsfull" -- Compiling hopes #96
-  ,"map" -- Compiling hopes #96
-  ,"infer_thunks" -- Weird: Mismatch between caller and callee signatures in map call
-  ,"infer_thunks2" -- Weird: Mismatch between caller and callee signatures in map call
   ,"repeated_app" -- missing coercions, https://github.com/quantinuum-dev/brat/issues/413
   ,"thunks"]
 
@@ -39,8 +35,7 @@ nonCompilingExamples = expectedCheckingFails ++ expectedParsingFails ++
   ,"let"
   ,"patterns"
   ,"qft"
-  ,"infer" -- problems with undoing pattern tests
-  ,"infer2" -- problems with undoing pattern tests
+  ,"infer2" -- https://github.com/Quantinuum/brat/issues/94
   ,"fanout" -- Contains Selectors
   ,"vectorise" -- Generates MapFun nodes which aren't implemented yet
   ,"vector_solve" -- Generates "Pow" nodes which aren't implemented yet
@@ -57,6 +52,11 @@ nonCompilingExamples = expectedCheckingFails ++ expectedParsingFails ++
   ,"vlup_covering"
   ]
 
+nonCompilingTests = map ((++ ".brat") . ("test" </>) . ("compilation" </>))
+  ["closures" -- https://github.com/Quantinuum/brat/issues/94
+  ,"parity" -- https://github.com/Quantinuum/brat/issues/94
+  ]
+
 compileToOutput :: FilePath -> TestTree
 compileToOutput file = testCaseInfo (show file) $ compileFile [] file >>= \case
     Right bs -> do
@@ -71,7 +71,7 @@ setupCompilationTests = do
   tests <- findByExtension [".brat"] prefix
   examples <- findByExtension [".brat"] examplesPrefix
   createDirectoryIfMissing False outputDir
-  let compileTests = compileToOutput <$> tests
+  let compileTests = expectFailForPaths nonCompilingTests compileToOutput tests
   let examplesTests = testGroup "examples" $ expectFailForPaths nonCompilingExamples compileToOutput examples
 
   pure $ testGroup "compilation" (examplesTests:compileTests)

brat/test/golden/graph/addN.brat.graph

@@ -6,7 +6,7 @@ Nodes:
 (check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs/in_1,BratNode Source [] [("n",Int)])
 (check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs/out_2,BratNode Target [("out",Int)] [])
 (check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs_thunk_3,BratNode (Box check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs/in_1 check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs/out_2) [] [("thunk",{ (n :: Int) -> (out :: Int) })])
-(check_defs_1_addN_LambdaChk_9_lambda,BratNode (PatternMatch ((TestMatchData Braty (MatchSequence {matchInputs = [(NamedPort {end = Ex check_defs_1_addN_LambdaChk_9_checkClauses_1_$lhs_lambda.0_setup/in 0, portName = "inp"},Int)], matchTests = [], matchOutputs = [(NamedPort {end = Ex check_defs_1_addN_LambdaChk_9_checkClauses_1_$lhs_lambda.0_setup/in 0, portName = "inp"},Int)]}),check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs_thunk_3) :| [])) [("inp",Int)] [("out",Int)])
+(check_defs_1_addN_LambdaChk_9_lambda,BratNode (PatternMatch ((TestMatchData Braty (MatchSequence {matchInputs = [(NamedPort {end = Ex check_defs_1_addN_LambdaChk_9_checkClauses_1_$lhs_lambda.0_setup/in 0, portName = "inp"},Int)], matchTests = [], rhsInputs = [(NamedPort {end = Ex check_defs_1_addN_LambdaChk_9_checkClauses_1_$lhs_lambda.0_setup/in 0, portName = "inp"},Int)]}),check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs_thunk_3) :| [])) [("inp",Int)] [("out",Int)])
 (check_defs_1_addN_addN.box/in,BratNode Source [] [("inp",Int)])
 (check_defs_1_addN_addN.box/out_1,BratNode Target [("out",Int)] [])
 (check_defs_1_addN_addN.box_thunk_2,BratNode (Box check_defs_1_addN_addN.box/in check_defs_1_addN_addN.box/out_1) [] [("thunk",{ (inp :: Int) -> (out :: Int) })])

brat/test/golden/graph/addN2.brat.graph

@@ -6,7 +6,7 @@ Nodes:
 (check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs/in_1,BratNode Source [] [("n",Int)])
 (check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs/out_2,BratNode Target [("out",Int)] [])
 (check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs_thunk_3,BratNode (Box check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs/in_1 check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs/out_2) [] [("thunk",{ (n :: Int) -> (out :: Int) })])
-(check_defs_1_addN_LambdaChk_9_lambda,BratNode (PatternMatch ((TestMatchData Braty (MatchSequence {matchInputs = [(NamedPort {end = Ex check_defs_1_addN_LambdaChk_9_checkClauses_1_$lhs_lambda.0_setup/in 0, portName = "inp"},Int)], matchTests = [], matchOutputs = [(NamedPort {end = Ex check_defs_1_addN_LambdaChk_9_checkClauses_1_$lhs_lambda.0_setup/in 0, portName = "inp"},Int)]}),check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs_thunk_3) :| [])) [("inp",Int)] [("out",Int)])
+(check_defs_1_addN_LambdaChk_9_lambda,BratNode (PatternMatch ((TestMatchData Braty (MatchSequence {matchInputs = [(NamedPort {end = Ex check_defs_1_addN_LambdaChk_9_checkClauses_1_$lhs_lambda.0_setup/in 0, portName = "inp"},Int)], matchTests = [], rhsInputs = [(NamedPort {end = Ex check_defs_1_addN_LambdaChk_9_checkClauses_1_$lhs_lambda.0_setup/in 0, portName = "inp"},Int)]}),check_defs_1_addN_LambdaChk_9_checkClauses_1_lambda.0_rhs_thunk_3) :| [])) [("inp",Int)] [("out",Int)])
 (check_defs_1_addN_addN.box/in,BratNode Source [] [("inp",Int)])
 (check_defs_1_addN_addN.box/out_1,BratNode Target [("out",Int)] [])
 (check_defs_1_addN_addN.box_thunk_2,BratNode (Box check_defs_1_addN_addN.box/in check_defs_1_addN_addN.box/out_1) [] [("thunk",{ (inp :: Int) -> (out :: Int) })])

brat/test/golden/graph/id.brat.graph

@@ -5,7 +5,7 @@ Nodes:
 (check_defs_1_main_$rhs_check'Th_LambdaChk_6_checkClauses_1_lambda.0_rhs/in_1,KernelNode Source [] [("q",Qubit)])
 (check_defs_1_main_$rhs_check'Th_LambdaChk_6_checkClauses_1_lambda.0_rhs/out_2,KernelNode Target [("b",Qubit)] [])
 (check_defs_1_main_$rhs_check'Th_LambdaChk_6_checkClauses_1_lambda.0_rhs_thunk_3,BratNode (Box check_defs_1_main_$rhs_check'Th_LambdaChk_6_checkClauses_1_lambda.0_rhs/in_1 check_defs_1_main_$rhs_check'Th_LambdaChk_6_checkClauses_1_lambda.0_rhs/out_2) [] [("thunk",{ (q :: Qubit) -o (b :: Qubit) })])
-(check_defs_1_main_$rhs_check'Th_LambdaChk_6_lambda,KernelNode (PatternMatch ((TestMatchData Kerny (MatchSequence {matchInputs = [(NamedPort {end = Ex check_defs_1_main_$rhs_check'Th_LambdaChk_6_checkClauses_1_$lhs_lambda.0_setup/in 0, portName = "a"},Qubit)], matchTests = [], matchOutputs = [(NamedPort {end = Ex check_defs_1_main_$rhs_check'Th_LambdaChk_6_checkClauses_1_$lhs_lambda.0_setup/in 0, portName = "a"},Qubit)]}),check_defs_1_main_$rhs_check'Th_LambdaChk_6_checkClauses_1_lambda.0_rhs_thunk_3) :| [])) [("a",Qubit)] [("b",Qubit)])
+(check_defs_1_main_$rhs_check'Th_LambdaChk_6_lambda,KernelNode (PatternMatch ((TestMatchData Kerny (MatchSequence {matchInputs = [(NamedPort {end = Ex check_defs_1_main_$rhs_check'Th_LambdaChk_6_checkClauses_1_$lhs_lambda.0_setup/in 0, portName = "a"},Qubit)], matchTests = [], rhsInputs = [(NamedPort {end = Ex check_defs_1_main_$rhs_check'Th_LambdaChk_6_checkClauses_1_$lhs_lambda.0_setup/in 0, portName = "a"},Qubit)]}),check_defs_1_main_$rhs_check'Th_LambdaChk_6_checkClauses_1_lambda.0_rhs_thunk_3) :| [])) [("a",Qubit)] [("b",Qubit)])
 (check_defs_1_main_$rhs_check'Th_thunk/in,KernelNode Source [] [("a",Qubit)])
 (check_defs_1_main_$rhs_check'Th_thunk/out_1,KernelNode Target [("b",Qubit)] [])
 (check_defs_1_main_$rhs_check'Th_thunk_thunk_2,BratNode (Box check_defs_1_main_$rhs_check'Th_thunk/in check_defs_1_main_$rhs_check'Th_thunk/out_1) [] [("thunk",{ (a :: Qubit) -o (b :: Qubit) })])