{-# OPTIONS_GHC -F -pgmF htfpp #-}
{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE TemplateHaskell #-}

{-
 - Generate a docker-compose compose.yml file
 -}

module Striot.CompileIoT.Compose
  ( generateDockerCompose
  ) where

import Striot.StreamGraph
import Striot.CompileIoT
import Algebra.Graph
import Data.List (intersect, null)
import Test.Framework
import Data.Maybe (fromJust, listToMaybe)

-- | Generate a docker-compose-format compose.yml in a String, encoding the
-- Node inter-connections from the supplied PartitionedGraph as dependencies,
-- such that consumer Nodes will be started prior to the producer Nodes that
-- connect to them.
generateDockerCompose :: PartitionedGraph -> String
generateDockerCompose :: PartitionedGraph -> String
generateDockerCompose pg :: PartitionedGraph
pg@([StreamGraph]
sgs,StreamGraph
_) = String
"services:\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++
    (Int -> String) -> [Int] -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (PartitionedGraph -> Int -> String
nodeToCompose PartitionedGraph
pg) [Int
1 .. ([StreamGraph] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [StreamGraph]
sgs)]

-- | Generate a YAML snippet in a String corresponding to the docker-compose
-- stanza for the Node in the PartitionedGraph at index i (using 1-indexing).
nodeToCompose :: PartitionedGraph -> Int -> String
nodeToCompose :: PartitionedGraph -> Int -> String
nodeToCompose PartitionedGraph
pg Int
i = [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
  [ String
"    ", String
n, String
":\n",
    String
"        build: ", String
n, String
"\n",
    String
"        tty: true\n",
    String
deps
  ] where
    n :: String
n = String
"node" String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Int -> String
forall a. Show a => a -> String
show Int
i)
    deps :: String
deps = case PartitionedGraph -> Int -> Maybe Int
getConsumer PartitionedGraph
pg Int
i of
      Maybe Int
Nothing -> String
""
      Just Int
d  -> String
"        depends_on:\n"
              String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"        - node" String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Int -> String
forall a. Show a => a -> String
show Int
d) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"\n"

-- | Return the index of the 'StreamGraph' downstream from that
-- at the given index, or 'Nothing' if there isn't one.
getConsumer :: PartitionedGraph -> Int -> Maybe Int
getConsumer :: PartitionedGraph -> Int -> Maybe Int
getConsumer ([StreamGraph]
sgs,StreamGraph
cuts) Int
i = let
  i' :: Int
i' = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 -- i is 1-indexed; we need 0-indexed
  vs :: [StreamVertex]
vs = StreamGraph -> [StreamVertex]
forall a. Ord a => Graph a -> [a]
vertexList ([StreamGraph]
sgs [StreamGraph] -> Int -> StreamGraph
forall a. [a] -> Int -> a
!! Int
i')
  receivingVertices :: [StreamVertex]
receivingVertices =
    (((StreamVertex, StreamVertex) -> StreamVertex)
-> [(StreamVertex, StreamVertex)] -> [StreamVertex]
forall a b. (a -> b) -> [a] -> [b]
map (StreamVertex, StreamVertex) -> StreamVertex
forall a b. (a, b) -> b
snd ([(StreamVertex, StreamVertex)] -> [StreamVertex])
-> (StreamGraph -> [(StreamVertex, StreamVertex)])
-> StreamGraph
-> [StreamVertex]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((StreamVertex, StreamVertex) -> Bool)
-> [(StreamVertex, StreamVertex)] -> [(StreamVertex, StreamVertex)]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(StreamVertex
a,StreamVertex
b) -> StreamVertex
a StreamVertex -> [StreamVertex] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [StreamVertex]
vs) ([(StreamVertex, StreamVertex)] -> [(StreamVertex, StreamVertex)])
-> (StreamGraph -> [(StreamVertex, StreamVertex)])
-> StreamGraph
-> [(StreamVertex, StreamVertex)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StreamGraph -> [(StreamVertex, StreamVertex)]
forall a. Ord a => Graph a -> [(a, a)]
edgeList) StreamGraph
cuts
  nsgs :: [(Int, StreamGraph)]
nsgs = [Int] -> [StreamGraph] -> [(Int, StreamGraph)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1..] [StreamGraph]
sgs
  res :: [(Int, StreamGraph)]
res = ((Int, StreamGraph) -> Bool)
-> [(Int, StreamGraph)] -> [(Int, StreamGraph)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool)
-> ((Int, StreamGraph) -> Bool) -> (Int, StreamGraph) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [StreamVertex] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([StreamVertex] -> Bool)
-> ((Int, StreamGraph) -> [StreamVertex])
-> (Int, StreamGraph)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [StreamVertex] -> [StreamVertex] -> [StreamVertex]
forall a. Eq a => [a] -> [a] -> [a]
intersect [StreamVertex]
receivingVertices ([StreamVertex] -> [StreamVertex])
-> ((Int, StreamGraph) -> [StreamVertex])
-> (Int, StreamGraph)
-> [StreamVertex]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StreamGraph -> [StreamVertex]
forall a. Ord a => Graph a -> [a]
vertexList (StreamGraph -> [StreamVertex])
-> ((Int, StreamGraph) -> StreamGraph)
-> (Int, StreamGraph)
-> [StreamVertex]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, StreamGraph) -> StreamGraph
forall a b. (a, b) -> b
snd) [(Int, StreamGraph)]
nsgs
  in [Int] -> Maybe Int
forall a. [a] -> Maybe a
listToMaybe (((Int, StreamGraph) -> Int) -> [(Int, StreamGraph)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, StreamGraph) -> Int
forall a b. (a, b) -> a
fst [(Int, StreamGraph)]
res)

-- test data, derived from examples/merge
v1 :: StreamVertex
v1 = Int
-> StreamOperator
-> [ExpQ]
-> String
-> String
-> Double
-> StreamVertex
StreamVertex Int
1 (Double -> StreamOperator
Source Double
1) [] String
"String" String
"String" Double
0
v2 :: StreamVertex
v2 = Int
-> StreamOperator
-> [ExpQ]
-> String
-> String
-> Double
-> StreamVertex
StreamVertex Int
2 (Double -> StreamOperator
Source Double
1) [] String
"String" String
"String" Double
0
v3 :: StreamVertex
v3 = Int
-> StreamOperator
-> [ExpQ]
-> String
-> String
-> Double
-> StreamVertex
StreamVertex Int
3 (Double -> StreamOperator
Source Double
1) [] String
"String" String
"String" Double
0
v4 :: StreamVertex
v4 = Int
-> StreamOperator
-> [ExpQ]
-> String
-> String
-> Double
-> StreamVertex
StreamVertex Int
4 StreamOperator
Merge [] String
"String" String
"String" Double
1
-- XXX: ^ we lie about the input type here, because the generated function has split-out arguments
v5 :: StreamVertex
v5 = Int
-> StreamOperator
-> [ExpQ]
-> String
-> String
-> Double
-> StreamVertex
StreamVertex Int
5 StreamOperator
Sink [[| mapM_ print |]] String
"String" String
"IO ()" Double
0

graph :: StreamGraph
graph = ([StreamGraph] -> StreamGraph
forall a. [Graph a] -> Graph a
overlays ((StreamVertex -> StreamGraph) -> [StreamVertex] -> [StreamGraph]
forall a b. (a -> b) -> [a] -> [b]
map StreamVertex -> StreamGraph
forall a. a -> Graph a
vertex [StreamVertex
v1,StreamVertex
v2,StreamVertex
v3]) StreamGraph -> StreamGraph -> StreamGraph
forall a. Graph a -> Graph a -> Graph a
`connect` (StreamVertex -> StreamGraph
forall a. a -> Graph a
vertex StreamVertex
v4)) StreamGraph -> StreamGraph -> StreamGraph
forall a. Graph a -> Graph a -> Graph a
`overlay` [StreamVertex] -> StreamGraph
forall a. [a] -> Graph a
path [StreamVertex
v4,StreamVertex
v5]
parts :: [[Int]]
parts = [[Int
1],[Int
2],[Int
3],[Int
4,Int
5]]
pg :: PartitionedGraph
pg    = StreamGraph -> [[Int]] -> PartitionedGraph
createPartitions StreamGraph
graph [[Int]]
parts

test_getConsumer_1 :: IO Int
test_getConsumer_1 = Maybe Int -> IO Int
forall a. HasCallStack => Maybe a -> IO a
assertJust (Maybe Int -> IO Int) -> Maybe Int -> IO Int
forall a b. (a -> b) -> a -> b
$ PartitionedGraph -> Int -> Maybe Int
getConsumer PartitionedGraph
pg Int
1
test_getConsumer_2 :: IO Int
test_getConsumer_2 = Maybe Int -> IO Int
forall a. HasCallStack => Maybe a -> IO a
assertJust (Maybe Int -> IO Int) -> Maybe Int -> IO Int
forall a b. (a -> b) -> a -> b
$ PartitionedGraph -> Int -> Maybe Int
getConsumer PartitionedGraph
pg Int
2
test_getConsumer_3 :: IO Int
test_getConsumer_3 = Maybe Int -> IO Int
forall a. HasCallStack => Maybe a -> IO a
assertJust (Maybe Int -> IO Int) -> Maybe Int -> IO Int
forall a b. (a -> b) -> a -> b
$ PartitionedGraph -> Int -> Maybe Int
getConsumer PartitionedGraph
pg Int
3
test_getConsumer_4 :: IO ()
test_getConsumer_4 = Maybe Int -> IO ()
forall a. (HasCallStack, Show a) => Maybe a -> IO ()
assertNothing (Maybe Int -> IO ()) -> Maybe Int -> IO ()
forall a b. (a -> b) -> a -> b
$ PartitionedGraph -> Int -> Maybe Int
getConsumer PartitionedGraph
pg Int
4

test_getConsumer_2a :: IO ()
test_getConsumer_2a = Maybe Int -> Maybe Int -> IO ()
forall a. (Eq a, Show a, HasCallStack) => a -> a -> IO ()
assertEqual (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
4) (Maybe Int -> IO ()) -> Maybe Int -> IO ()
forall a b. (a -> b) -> a -> b
$ PartitionedGraph -> Int -> Maybe Int
getConsumer PartitionedGraph
pg Int
2