|
| 1 | +{-# LANGUAGE FlexibleContexts #-} |
| 2 | +{-# LANGUAGE GADTSyntax #-} |
| 3 | +{-# LANGUAGE ImportQualifiedPost #-} |
| 4 | +{-# LANGUAGE LambdaCase #-} |
| 5 | +{-# LANGUAGE OverloadedStrings #-} |
| 6 | +{-# LANGUAGE ScopedTypeVariables #-} |
| 7 | +{-# LANGUAGE TupleSections #-} |
| 8 | + |
| 9 | +module Swarm.DocGen ( |
| 10 | + generateDocs, |
| 11 | + GenerateDocs (..), |
| 12 | +) where |
| 13 | + |
| 14 | +import Control.Lens (view, (^.)) |
| 15 | +import Control.Monad (zipWithM, zipWithM_, (<=<)) |
| 16 | +import Control.Monad.Except (ExceptT, runExceptT) |
| 17 | +import Data.Bifunctor (Bifunctor (bimap)) |
| 18 | +import Data.Containers.ListUtils (nubOrd) |
| 19 | +import Data.Foldable (toList) |
| 20 | +import Data.Map.Lazy (Map) |
| 21 | +import Data.Map.Lazy qualified as Map |
| 22 | +import Data.Maybe (fromMaybe) |
| 23 | +import Data.Set (Set) |
| 24 | +import Data.Set qualified as Set |
| 25 | +import Data.Text (Text, unpack) |
| 26 | +import Data.Tuple (swap) |
| 27 | +import Swarm.Game.Entity (Entity, EntityMap (entitiesByName), entityName, loadEntities) |
| 28 | +import Swarm.Game.Entity qualified as E |
| 29 | +import Swarm.Game.Recipe (Recipe, loadRecipes, recipeInputs, recipeOutputs, recipeRequirements) |
| 30 | +import Swarm.Game.Robot (installedDevices, robotInventory, setRobotID) |
| 31 | +import Swarm.Game.Scenario (Scenario, loadScenario, scenarioRobots) |
| 32 | +import Swarm.Game.WorldGen (testWorld2Entites) |
| 33 | +import Swarm.Util (isRightOr) |
| 34 | +import Text.Dot (Dot, NodeId, (.->.)) |
| 35 | +import Text.Dot qualified as Dot |
| 36 | + |
| 37 | +-- ============================================================================ |
| 38 | +-- MAIN ENTRYPOINT TO CLI DOCUMENTATION GENERATOR |
| 39 | +-- ============================================================================ |
| 40 | +-- |
| 41 | +-- These are the exported functions used by the executable. |
| 42 | +-- |
| 43 | +-- ---------------------------------------------------------------------------- |
| 44 | + |
| 45 | +data GenerateDocs where |
| 46 | + -- | Entity dependencies by recipes. |
| 47 | + RecipeGraph :: GenerateDocs |
| 48 | + deriving (Eq, Show) |
| 49 | + |
| 50 | +generateDocs :: GenerateDocs -> IO () |
| 51 | +generateDocs = \case |
| 52 | + RecipeGraph -> generateRecipe >>= putStrLn |
| 53 | + |
| 54 | +-- ---------------------------------------------------------------------------- |
| 55 | +-- GENERATE GRAPHVIZ: ENTITY DEPENDENCIES BY RECIPES |
| 56 | +-- ---------------------------------------------------------------------------- |
| 57 | + |
| 58 | +generateRecipe :: IO String |
| 59 | +generateRecipe = simpleErrorHandle $ do |
| 60 | + entities <- loadEntities >>= guardRight "load entities" |
| 61 | + recipes <- loadRecipes entities >>= guardRight "load recipes" |
| 62 | + classic <- classicScenario |
| 63 | + return . Dot.showDot $ recipesToDot classic entities recipes |
| 64 | + |
| 65 | +recipesToDot :: Scenario -> EntityMap -> [Recipe Entity] -> Dot () |
| 66 | +recipesToDot classic emap recipes = do |
| 67 | + Dot.attribute ("rankdir", "LR") |
| 68 | + Dot.attribute ("ranksep", "2") |
| 69 | + world <- diamond "World" |
| 70 | + base <- diamond "Base" |
| 71 | + -- -------------------------------------------------------------------------- |
| 72 | + -- add nodes with for all the known entites |
| 73 | + let enames' = toList . Map.keysSet . entitiesByName $ emap |
| 74 | + enames = filter (`Set.notMember` ignoredEntites) enames' |
| 75 | + ebmap <- Map.fromList . zip enames <$> mapM (box . unpack) enames |
| 76 | + -- -------------------------------------------------------------------------- |
| 77 | + -- getters for the NodeId based on entity name or the whole entity |
| 78 | + let safeGetEntity m e = fromMaybe (error $ unpack e <> " is not an entity!?") $ m Map.!? e |
| 79 | + getE = safeGetEntity ebmap |
| 80 | + nid = getE . view entityName |
| 81 | + -- -------------------------------------------------------------------------- |
| 82 | + -- Get the starting inventories, entites present in the world and compute |
| 83 | + -- how hard each entity is to get - see 'recipeLevels'. |
| 84 | + let devs = startingDevices classic |
| 85 | + inv = startingInventory classic |
| 86 | + worldEntites = Set.map (safeGetEntity $ entitiesByName emap) testWorld2Entites |
| 87 | + levels = recipeLevels recipes (Set.unions [worldEntites, devs]) |
| 88 | + -- -------------------------------------------------------------------------- |
| 89 | + -- Base inventory |
| 90 | + (_bc, ()) <- Dot.cluster $ do |
| 91 | + Dot.attribute ("style", "filled") |
| 92 | + Dot.attribute ("color", "lightgrey") |
| 93 | + mapM_ ((base ---<>) . nid) devs |
| 94 | + mapM_ ((base .->.) . nid . fst) $ Map.toList inv |
| 95 | + -- -------------------------------------------------------------------------- |
| 96 | + -- World entites |
| 97 | + (_wc, ()) <- Dot.cluster $ do |
| 98 | + Dot.attribute ("style", "filled") |
| 99 | + Dot.attribute ("color", "forestgreen") |
| 100 | + mapM_ ((uncurry (Dot..->.) . (world,)) . getE) (toList testWorld2Entites) |
| 101 | + -- -------------------------------------------------------------------------- |
| 102 | + let -- put a hidden node above and below entites and connect them by hidden edges |
| 103 | + wrapBelowAbove :: Set Entity -> Dot (NodeId, NodeId) |
| 104 | + wrapBelowAbove ns = do |
| 105 | + b <- hiddenNode |
| 106 | + t <- hiddenNode |
| 107 | + let ns' = map nid $ toList ns |
| 108 | + mapM_ (b .~>.) ns' |
| 109 | + mapM_ (.~>. t) ns' |
| 110 | + return (b, t) |
| 111 | + -- put set of entites in nice |
| 112 | + subLevel :: Int -> Set Entity -> Dot (NodeId, NodeId) |
| 113 | + subLevel i ns = fmap snd . Dot.cluster $ do |
| 114 | + Dot.attribute ("style", "filled") |
| 115 | + Dot.attribute ("color", "khaki") |
| 116 | + bt <- wrapBelowAbove ns |
| 117 | + Dot.attribute ("rank", "sink") |
| 118 | + -- the normal label for cluster would be cover by lines |
| 119 | + _bigLabel <- |
| 120 | + Dot.node |
| 121 | + [ ("shape", "plain") |
| 122 | + , ("label", "Bottom Label") |
| 123 | + , ("fontsize", "20pt") |
| 124 | + , ("label", "Level #" <> show i) |
| 125 | + ] |
| 126 | + return bt |
| 127 | + -- -------------------------------------------------------------------------- |
| 128 | + -- order entites into clusters based on how "far" they are from |
| 129 | + -- what is available at the start - see 'recipeLevels'. |
| 130 | + bottom <- wrapBelowAbove worldEntites |
| 131 | + ls <- zipWithM subLevel [1 ..] (tail levels) |
| 132 | + let invisibleLine = zipWithM_ (.~>.) |
| 133 | + tls <- mapM (const hiddenNode) levels |
| 134 | + bls <- mapM (const hiddenNode) levels |
| 135 | + invisibleLine tls bls |
| 136 | + invisibleLine bls (tail tls) |
| 137 | + let sameBelowAbove (b1, t1) (b2, t2) = Dot.same [b1, b2] >> Dot.same [t1, t2] |
| 138 | + zipWithM_ sameBelowAbove (bottom : ls) (zip bls tls) |
| 139 | + -- -------------------------------------------------------------------------- |
| 140 | + -- add node for the world and draw a line to each entity found in the wild |
| 141 | + -- finally draw recipes |
| 142 | + let recipeInOut r = [(snd i, snd o) | i <- r ^. recipeInputs, o <- r ^. recipeOutputs] |
| 143 | + recipeReqOut r = [(snd q, snd o) | q <- r ^. recipeRequirements, o <- r ^. recipeOutputs] |
| 144 | + recipesToPairs f rs = both nid <$> nubOrd (concatMap f rs) |
| 145 | + mapM_ (uncurry (.->.)) (recipesToPairs recipeInOut recipes) |
| 146 | + mapM_ (uncurry (---<>)) (recipesToPairs recipeReqOut recipes) |
| 147 | + |
| 148 | +-- ---------------------------------------------------------------------------- |
| 149 | +-- RECIPE LEVELS |
| 150 | +-- ---------------------------------------------------------------------------- |
| 151 | + |
| 152 | +-- | Order entites in sets depending on how soon it is possible to obtain them. |
| 153 | +-- |
| 154 | +-- So: |
| 155 | +-- * Level 0 - starting entites (for example those obtainable in the world) |
| 156 | +-- * Level N+1 - everything possible to make (or drill) from Level N |
| 157 | +-- |
| 158 | +-- This is almost a BFS, but the requirement is that the set of entites |
| 159 | +-- required for recipe is subset of the entites known in Level N. |
| 160 | +-- |
| 161 | +-- If we ever depend on some graph library, this could be rewritten |
| 162 | +-- as some BFS-like algorithm with added recipe nodes, but you would |
| 163 | +-- need to enforce the condition that recipes need ALL incoming edges. |
| 164 | +recipeLevels :: [Recipe Entity] -> Set Entity -> [Set Entity] |
| 165 | +recipeLevels recipes start = levels |
| 166 | + where |
| 167 | + recipeParts r = ((r ^. recipeInputs) <> (r ^. recipeRequirements), r ^. recipeOutputs) |
| 168 | + m :: [(Set Entity, Set Entity)] |
| 169 | + m = map (both (Set.fromList . map snd) . recipeParts) recipes |
| 170 | + levels :: [Set Entity] |
| 171 | + levels = reverse $ go [start] start |
| 172 | + where |
| 173 | + isKnown known (i, _o) = null $ i Set.\\ known |
| 174 | + nextLevel known = Set.unions . map snd $ filter (isKnown known) m |
| 175 | + go ls known = |
| 176 | + let n = nextLevel known Set.\\ known |
| 177 | + in if null n |
| 178 | + then ls |
| 179 | + else go (n : ls) (Set.union n known) |
| 180 | + |
| 181 | +-- | Get classic scenario to figure out starting entites. |
| 182 | +classicScenario :: ExceptT Text IO Scenario |
| 183 | +classicScenario = do |
| 184 | + entities <- loadEntities >>= guardRight "load entities" |
| 185 | + loadScenario "data/scenarios/00-classic.yaml" entities |
| 186 | + |
| 187 | +startingDevices :: Scenario -> Set Entity |
| 188 | +startingDevices = Set.fromList . map snd . E.elems . view installedDevices . setRobotID 0 . head . view scenarioRobots |
| 189 | + |
| 190 | +startingInventory :: Scenario -> Map Entity Int |
| 191 | +startingInventory = Map.fromList . map swap . E.elems . view robotInventory . setRobotID 0 . head . view scenarioRobots |
| 192 | + |
| 193 | +-- | Ignore utility entites that are just used for tutorials and challenges. |
| 194 | +ignoredEntites :: Set Text |
| 195 | +ignoredEntites = |
| 196 | + Set.fromList |
| 197 | + [ "upper left corner" |
| 198 | + , "upper right corner" |
| 199 | + , "lower left corner" |
| 200 | + , "lower right corner" |
| 201 | + , "horizontal wall" |
| 202 | + , "vertical wall" |
| 203 | + ] |
| 204 | + |
| 205 | +-- ---------------------------------------------------------------------------- |
| 206 | +-- GRAPHVIZ HELPERS |
| 207 | +-- ---------------------------------------------------------------------------- |
| 208 | + |
| 209 | +customNode :: [(String, String)] -> String -> Dot NodeId |
| 210 | +customNode attrs label = Dot.node $ [("style", "filled"), ("label", label)] <> attrs |
| 211 | + |
| 212 | +box, diamond :: String -> Dot NodeId |
| 213 | +box = customNode [("shape", "box")] |
| 214 | +diamond = customNode [("shape", "diamond")] |
| 215 | + |
| 216 | +-- | Hidden node - used for layout. |
| 217 | +hiddenNode :: Dot NodeId |
| 218 | +hiddenNode = Dot.node [("style", "invis")] |
| 219 | + |
| 220 | +-- | Hidden edge - used for layout. |
| 221 | +(.~>.) :: NodeId -> NodeId -> Dot () |
| 222 | +i .~>. j = Dot.edge i j [("style", "invis")] |
| 223 | + |
| 224 | +-- | Edge for recipe requirements and outputs. |
| 225 | +(---<>) :: NodeId -> NodeId -> Dot () |
| 226 | +e1 ---<> e2 = Dot.edge e1 e2 attrs |
| 227 | + where |
| 228 | + attrs = [("arrowhead", "diamond"), ("color", "blue")] |
| 229 | + |
| 230 | +-- ---------------------------------------------------------------------------- |
| 231 | +-- UTILITY |
| 232 | +-- ---------------------------------------------------------------------------- |
| 233 | + |
| 234 | +both :: Bifunctor p => (a -> d) -> p a a -> p d d |
| 235 | +both f = bimap f f |
| 236 | + |
| 237 | +guardRight :: Text -> Either Text a -> ExceptT Text IO a |
| 238 | +guardRight what i = i `isRightOr` (\e -> "Failed to " <> what <> ": " <> e) |
| 239 | + |
| 240 | +simpleErrorHandle :: ExceptT Text IO a -> IO a |
| 241 | +simpleErrorHandle = either (fail . unpack) pure <=< runExceptT |
0 commit comments