make scatter partition payloads using routed trees

Author: shayne-fletcher

src/Tile/Execution.hs

@@ -11,6 +11,7 @@ import Control.Monad
 import Data.Map.Strict qualified as Map
 import Data.Set qualified as Set
 import Tile.Schedule
+import Tile.Tree (RoutedTree (..), scheduleTree, treeIndex, treeLabels)
 
 runBroadcast :: Schedule String -> String -> IO ()
 runBroadcast schedule root = do
@@ -147,7 +148,8 @@ runScatter schedule initialValues root = do
         Set.toList $
           Set.fromList $
             Map.keys graph ++ concat (Map.elems graph) ++ map fst initialValues
-      subtree = subtreeMembers graph
+      RoutedTree routed = scheduleTree root schedule
+      routedSubtrees = treeIndex routed
 
   chanPairs <- forM members $ \m -> do
     ch <- newChan
@@ -164,13 +166,17 @@ runScatter schedule initialValues root = do
           | otherwise = Map.findWithDefault 0 m incoming
 
     _ <- forkIO $ do
+      -- Scatter schedules are normally trees, so this usually reads
+      -- one payload. For a general schedule, merge all incoming
+      -- payload fragments before forwarding.
       payload <- concat <$> replicateM expected (readChan inbox)
       case lookup m payload of
         Just value -> putStrLn $ m ++ " received scatter value: " ++ show value
         Nothing -> pure ()
 
       forM_ childChans $ \(childName, childInbox) -> do
-        let childMembers = subtree childName
+        let childMembers =
+              maybe Set.empty treeLabels (Map.lookup childName routedSubtrees)
             childPayload =
               [ item
               | item@(dest, _) <- payload,
@@ -182,11 +188,3 @@ runScatter schedule initialValues root = do
 
   writeChan (chanMap Map.! root) initialValues
   threadDelay 1000000
-
-subtreeMembers :: Map.Map String [String] -> String -> Set.Set String
-subtreeMembers graph member =
-  Set.insert member $
-    Set.unions
-      [ subtreeMembers graph child
-      | child <- Map.findWithDefault [] member graph
-      ]

src/Tile/Tree.hs

@@ -17,6 +17,8 @@ module Tile.Tree
     Tree (..),
     mapTree,
     unfoldTree,
+    treeLabels,
+    treeIndex,
     renderTreeWith,
 
     -- * Tile tree views
@@ -49,6 +51,7 @@ where
 
 import Data.List (sortOn)
 import Data.Map.Strict qualified as Map
+import Data.Set qualified as Set
 import Tile.Geometry
 import Tile.Schedule
 import Tile.Tile
@@ -210,6 +213,16 @@ unfoldTree childrenOf label =
       subtrees = map (unfoldTree childrenOf) (childrenOf label)
     }
 
+-- | Collect every label in a tree.
+treeLabels :: (Ord a) => Tree a -> Set.Set a
+treeLabels (Tree label kids) =
+  Set.insert label (Set.unions (map treeLabels kids))
+
+-- | Index every subtree by its root label.
+treeIndex :: (Ord a) => Tree a -> Map.Map a (Tree a)
+treeIndex tree@(Tree label kids) =
+  Map.insert label tree (Map.unions (map treeIndex kids))
+
 -- | Render a tree as an ASCII box-drawing string.
 renderTreeWith :: (a -> String) -> Tree a -> String
 renderTreeWith renderLabel tree =

test/Main.hs

@@ -1,6 +1,8 @@
 module Main (main) where
 
 import Data.List (sort)
+import Data.Map.Strict qualified as Map
+import Data.Set qualified as Set
 import Test.Tasty
 import Test.Tasty.HUnit
 import Test.Tasty.QuickCheck
@@ -288,7 +290,11 @@ treeTests =
                         ]
                     }
                 ]
-            }
+            },
+      testCase "treeLabels collects subtree labels" $
+        treeLabels sampleTree @?= Set.fromList [0, 1, 2, 3],
+      testCase "treeIndex indexes every subtree by label" $
+        Map.keysSet (treeIndex sampleTree) @?= Set.fromList [0, 1, 2, 3]
     ]
   where
     sampleTree :: Tree Int

tile.cabal

@@ -66,6 +66,7 @@ test-suite tile-test
     main-is:          Main.hs
     build-depends:
         base ^>=4.22.0.0,
+        containers >=0.8 && <0.9,
         tile,
         tasty >=1.4,
         tasty-hunit >=0.10,