examples/GArrowTikZ.hs

   1 {-# LANGUAGE RankNTypes, MultiParamTypeClasses, GADTs, FlexibleContexts, FlexibleInstances, TypeOperators #-}
   2 -----------------------------------------------------------------------------
   3 -- |
   4 -- Module      :  GArrowTikZ
   5 -- Copyright   :  none
   6 -- License     :  public domain
   7 --
   8 -- Maintainer  :  Adam Megacz <megacz@acm.org>
   9 -- Stability   :  experimental
  10 --
  11 -- | Renders a @GArrowSkeleton@ using TikZ; the result is LaTeX code
  12 --
  13
  14 module GArrowTikZ (tikz, tikz')
  15 where
  16 import Prelude hiding ( id, (.), lookup )
  17 import Control.Category
  18 import GHC.HetMet.GArrow
  19 import Data.List hiding (lookup, insert)
  20 import Data.Map hiding (map, (!))
  21 import Unify
  22 import GArrowSkeleton
  23 import GHC.HetMet.Private
  24
  25 data UPort =
  26    UPortVar  UVar
  27  | UPortPair UPort UPort
  28
  29 instance Unifiable UPort where
  30   unify' (UPortPair vl1a vl1b) (UPortPair vl2a vl2b) = mergeU (unify vl1a vl2a) (unify vl1b vl2b)
  31   unify' _ _                                         = error "impossible"
  32   inject                                             = UPortVar
  33   project (UPortVar v)                               = Just v
  34   project _                                          = Nothing
  35   occurrences (UPortVar v)                           = [v]
  36   occurrences (UPortPair x y)                        = occurrences x ++ occurrences y
  37
  38 -- | Resolves a unification variable; the answer to this query will change if subsequent unifications are performed
  39 getU' :: Unifier UPort -> UPort -> UPort
  40 getU' u (UPortPair x y)  = UPortPair (getU' u x) (getU' u y)
  41 getU' u x@(UPortVar v)   = case Unify.getU u v  of
  42                                      Nothing -> x
  43                                      Just x' -> getU' u x'
  44
  45
  46
  47
  48 --
  49 -- | Render a fully-polymorphic GArrow term as a boxes-and-wires diagram using TikZ
  50 --
  51
  52 type Constraints = [(UPort, Int, UPort)]
  53
  54 -- the unification monad
  55 data UyM t a = UyM (([UVar],Unifier UPort,Constraints) -> ([UVar],Unifier UPort,Constraints,a))
  56 instance Monad (UyM t)
  57  where
  58   return x      = UyM $ \(i,u,k) -> (i,u,k,x)
  59   (UyM f) >>= g = UyM $ \(i,u,k) -> let (i',u',k',x) = f (i,u,k) in let UyM g' = g x in g' (i',u',k')
  60
  61
  62 getU        = UyM $ \(i,u,k) -> (i,u,k,u)
  63 getM    v   = UyM $ \(i,u,k) -> (i,u,k,getU' u v)
  64 occursU v x = UyM $ \(i,u,k) -> (i,u,k,occurs u v x)
  65 unifyM :: Eq t => UPort -> UPort -> UyM t ()
  66 unifyM v1 v2 = UyM $ \(i,u,k) -> (i,mergeU u (unify v1 v2),k,())
  67 freshU :: UyM t UVar
  68 freshU = UyM $ \(i:is,u,k) -> (is,u,k,i)
  69 constrain :: UPort -> Int -> UPort -> UyM t ()
  70 constrain v1 d v2 = UyM $ \(i,u,k) -> (i,u,((v1,d,v2):k),())
  71 getK :: UyM t [(UPort, Int, UPort)]
  72 getK = UyM $ \(i,u,k) -> (i,u,k,k)
  73 runU :: UyM t a -> ([UVar],Unifier UPort,Constraints,a)
  74 runU (UyM f) = (f (uvarSupply,emptyUnifier,[]))
  75
  76
  77
  78 name :: GArrowSkeleton m a b -> String
  79 name GAS_id              = "id"
  80 name (GAS_const i)       = "const " ++ show i
  81 name (GAS_comp      _ _) = "comp"
  82 name (GAS_first     _  ) = "first"
  83 name (GAS_second    _  ) = "second"
  84 name GAS_cancell         = "cancell"
  85 name GAS_cancelr         = "cancelr"
  86 name GAS_uncancell       = "uncancell"
  87 name GAS_uncancelr       = "uncancelr"
  88 name GAS_drop            = "drop"
  89 name GAS_copy            = "copy"
  90 name GAS_swap            = "swap"
  91 name (GAS_loopl     _ )  = "loopl"
  92 name (GAS_loopr     _ )  = "loopr"
  93 name GAS_assoc           = "assoc"
  94 name GAS_unassoc         = "unassoc"
  95 name GAS_merge           = "merge"
  96 name (GAS_misc _)        = "misc"
  97
  98 fresh :: Int -> UyM () [Ports]
  99 fresh 0 = return []
 100 fresh n = do { xs <- fresh  (n-1)
 101              ; x  <- freshU
 102              ; case xs of
 103                  []       -> return [UPortVar x]
 104                  (x':xs') -> do { constrain (UPortVar x) 1 x'
 105                                 ; return $ (UPortVar x):x':xs'
 106                                 }
 107              }
 108
 109 data Diagram p = DiagramComp      (Diagram p) (Diagram p)
 110                | DiagramPrim      String p p p p (String -> Int -> Int -> Int -> p -> p -> Int -> String)
 111                | DiagramBypassTop p (Diagram p)
 112                | DiagramBypassBot (Diagram p) p
 113                -- | DiagramLoopTop   Diagram
 114                -- | DiagramLoopBot   Diagram
 115
 116 type Ports = UPort
 117
 118 getOut :: Diagram Ports -> Ports
 119 getOut (DiagramComp f g)                     = getOut g
 120 getOut (DiagramPrim s ptop pin pout pbot _)  = pout
 121 getOut (DiagramBypassTop p f)                = UPortPair p (getOut f)
 122 getOut (DiagramBypassBot f p)                = UPortPair (getOut f) p
 123
 124 getIn :: Diagram Ports -> Ports
 125 getIn (DiagramComp f g)                      = getIn f
 126 getIn (DiagramPrim s ptop pin pout pbot _)   = pin
 127 getIn (DiagramBypassTop p f)                 = UPortPair p (getIn f)
 128 getIn (DiagramBypassBot f p)                 = UPortPair (getIn f) p
 129
 130 -- constrain that Ports is at least Int units above the topmost portion of Diagram
 131 constrainTop :: Ports -> Int -> Diagram Ports -> UyM () ()
 132 constrainTop v i (DiagramComp d1 d2)                  = do { constrainTop v i d1 ; constrainTop v i d2 ; return () }
 133 constrainTop v i (DiagramBypassTop p d)               = constrain v 2 p
 134 constrainTop v i (DiagramBypassBot d p)               = constrainTop v (i+1) d
 135 constrainTop v i (DiagramPrim s ptop pin pout pbot _) = constrain v i ptop
 136
 137 -- constrain that Ports is at least Int units below the bottommost portion of Diagram
 138 constrainBot :: Diagram Ports -> Int -> Ports -> UyM () ()
 139 constrainBot (DiagramComp d1 d2)                  i v = do { constrainBot d1 i v ; constrainBot d2 i v ; return () }
 140 constrainBot (DiagramBypassTop p d)               i v = constrainBot d (i+1) v
 141 constrainBot (DiagramBypassBot d p)               i v = constrain p 2 v
 142 constrainBot (DiagramPrim s ptop pin pout pbot _) i v = constrain pbot i v
 143
 144 ga2diag :: GArrowSkeleton m a b -> UyM () (Diagram Ports)
 145 ga2diag (GAS_id           ) = do { [top,x,bot] <- fresh 3 ; return $ DiagramPrim "id" top x x bot defren }
 146 ga2diag (GAS_comp      f g) = do { f' <- ga2diag f
 147                                   ; g' <- ga2diag g
 148                                   ; unifyM (getOut f') (getIn g')
 149                                   ; return $ DiagramComp f' g' }
 150 ga2diag (GAS_first  f) = do { [x] <- fresh 1; f' <- ga2diag f ; constrainBot f' 1 x  ; return $ DiagramBypassBot f' x  }
 151 ga2diag (GAS_second f) = do { [x] <- fresh 1; f' <- ga2diag f ; constrainTop x  1 f' ; return $ DiagramBypassTop x f'  }
 152 ga2diag GAS_cancell    = do { [top,x,y  ,bot] <- fresh 4  ; return $ DiagramPrim   "cancell" top (UPortPair x y) y bot defren }
 153 ga2diag GAS_cancelr    = do { [top,x,y  ,bot] <- fresh 4  ; return $ DiagramPrim   "cancelr" top (UPortPair x y) x bot defren }
 154 ga2diag GAS_uncancell  = do { [top,x,y  ,bot] <- fresh 4  ; return $ DiagramPrim "uncancell" top y (UPortPair x y) bot defren }
 155 ga2diag GAS_uncancelr  = do { [top,x,y  ,bot] <- fresh 4  ; return $ DiagramPrim "uncancelr" top x (UPortPair x y) bot defren }
 156 ga2diag GAS_drop       = do { [top,x    ,bot] <- fresh 3  ; return $ DiagramPrim      "drop" top x x bot defren }
 157 ga2diag (GAS_const i)  = do { [top,x    ,bot] <- fresh 3  ; return $ DiagramPrim  ("const " ++ show i) top x x bot defren }
 158 ga2diag GAS_copy       = do { [top,x,y,z,bot] <- fresh 5
 159                              ; return $ DiagramPrim      "copy" top y (UPortPair x z) bot defren }
 160 ga2diag GAS_merge      = do { [top,x,y,z,bot] <- fresh 5
 161                              ; return $ DiagramPrim      "merge" top (UPortPair x z) y bot defren }
 162 ga2diag GAS_swap       = do { [top,x,y  ,bot] <- fresh 4
 163                              ; return $ DiagramPrim      "swap" top (UPortPair x y) (UPortPair x y) bot defren }
 164 ga2diag GAS_assoc      = do { [top,x,y,z,bot] <- fresh 5
 165                              ; return $ DiagramPrim  "assoc" top (UPortPair (UPortPair x y) z) (UPortPair x (UPortPair y z)) bot defren }
 166 ga2diag GAS_unassoc    = do { [top,x,y,z,bot] <- fresh 5
 167                              ; return $ DiagramPrim "unassoc" top (UPortPair x (UPortPair y z)) (UPortPair (UPortPair x y) z) bot defren }
 168 ga2diag (GAS_loopl f) = error "not implemented"
 169 ga2diag (GAS_loopr f) = error "not implemented"
 170 ga2diag (GAS_misc f ) = error "not implemented"
 171
 172 defren :: String -> Int -> Int -> Int -> Ports -> Ports -> Int -> String
 173 defren s x w top p1 p2 bot
 174       = drawBox x top (x+w) bot "black" s
 175 --        ++ wires (x-1) p1  x    "green"
 176 --        ++ wires  (x+w) p2 (x+w+1) "red"
 177
 178 xscale = 1
 179 yscale = 1
 180
 181 textc x y text color =
 182     "\\node[anchor=center,color="++color++"] at ("++show (x*xscale)++"cm,"++show (y*yscale)++"cm) "++
 183     "{{\\tt{"++text++"}}};\n"
 184
 185 drawBox x1 y1 x2 y2 color text =
 186     "\\node[anchor=north west] at ("++show (x1*xscale)++"cm,"++show (y1*yscale)++"cm) "++
 187     "{{\\tt{"++text++"}}};\n"
 188     ++
 189     "\\path[draw,color="++color++"]"++
 190        " ("++show (x1*xscale)++","++show (y1*yscale)++") rectangle ("++
 191            show (x2*xscale)++","++show (y2*yscale)++");\n"
 192
 193 drawLine x1 y1 x2 y2 color style =
 194   "\\path[draw,color="++color++","++style++"] "++
 195   "("++show (x1*xscale)++","++show (y1*yscale)++") -- " ++
 196   "("++show (x2*xscale)++","++show (y2*yscale)++");\n"
 197
 198 width :: Diagram Ports -> Int
 199 width (DiagramComp d1 d2)         = (width d1) + 1 + (width d2)
 200 width (DiagramPrim s _ p1 p2 _ _) = 2
 201 width (DiagramBypassTop p d)      = (width d) + 2
 202 width (DiagramBypassBot d p)      = (width d) + 2
 203
 204 (!) :: Map UVar Int -> UVar -> Int
 205 m ! x = case lookup x m of
 206           Nothing -> 0
 207           Just y -> y
 208
 209 tikZ :: Map UVar Int ->
 210         Diagram Ports ->
 211         Int ->                -- horizontal position
 212         String
 213 tikZ m = tikZ'
 214  where
 215   tikZ'  d@(DiagramComp d1 d2)    x = tikZ' d1 x
 216                                       ++ wires (x+width d1) (getOut d1) (x+width d1+1) "black"
 217                                       ++ tikZ' d2 (x + width d1 + 1)
 218   tikZ' d'@(DiagramBypassTop p d) x = let top = getTop d' in
 219                                       let bot = getBot d' in
 220                                       drawBox  x top (x+width d') bot "gray!50" "second"
 221                                       ++ wires x (getIn d) (x+1) "black"
 222                                       ++ tikZ' d (x+1)
 223                                       ++ wires (x+1+width d) (getOut d) (x+1+width d+1) "black"
 224                                       ++ wires x p (x+1+width d+1) "black"
 225   tikZ' d'@(DiagramBypassBot d p) x = let top = getTop d' in
 226                                       let bot = getBot d' in
 227                                       drawBox  x top (x+width d') bot "gray!50" "first"
 228                                       ++ wires x (getIn d) (x+1) "black"
 229                                       ++ tikZ' d (x+1)
 230                                       ++ wires (x+1+width d) (getOut d) (x+1+width d+1) "black"
 231                                       ++ wires x p (x+1+width d+1) "black"
 232   tikZ'  d@(DiagramPrim s (UPortVar top) p1 p2 (UPortVar bot) r)  x = r s x (width d) (m ! top) p1 p2 (m ! bot)
 233   tikZ'  _ _ = error "FIXME"
 234
 235   wires :: Int -> Ports -> Int -> String -> String
 236   wires x1 (UPortVar v)    x2 color = drawLine x1 (m ! v) x2 (m ! v) color "->"
 237                                        --  ++ textc ((x1+x2) `div` 2) (m!v) (show v) "purple"
 238   wires x1 (UPortPair x y) x2 color = wires x1 x x2 color ++ wires x1 y x2 color
 239
 240   getTop :: Diagram Ports -> Int
 241   getTop (DiagramComp d1 d2)                = min (getTop d1) (getTop d2)
 242   getTop (DiagramBypassTop p d)             = (m ! getleft p) - 1
 243   getTop (DiagramBypassBot d p)             = getTop d - 1
 244   getTop (DiagramPrim s (UPortVar ptop) _ _ _ _)  = m ! ptop
 245   getTop _ = error "getTop failed"
 246
 247   getBot :: Diagram Ports -> Int
 248   getBot (DiagramComp d1 d2)                = max (getBot d1) (getBot d2)
 249   getBot (DiagramBypassTop p d)             = getBot d + 1
 250   getBot (DiagramBypassBot d p)             = (m ! getright p) + 1
 251   getBot (DiagramPrim s _ _ _ (UPortVar pbot) _)  = m ! pbot
 252   getBot _ = error "getTop failed"
 253
 254 resolve' (DiagramComp d1 d2)    = do { d1' <- resolve' d1 ; d2' <- resolve' d2 ; return $ DiagramComp d1' d2'    }
 255 resolve' (DiagramBypassTop p d) = do { p'  <- getM p     ; d'  <- resolve' d  ; return $ DiagramBypassTop p' d' }
 256 resolve' (DiagramBypassBot d p) = do { p'  <- getM p     ; d'  <- resolve' d  ; return $ DiagramBypassBot d' p' }
 257 resolve' (DiagramPrim s ptop pin pout pbot r)
 258     = do { ptop' <- getM ptop
 259          ; pbot' <- getM pbot
 260          ; pin'  <- getM pin
 261          ; pout' <- getM pout
 262          ; return $ DiagramPrim s ptop' pin' pout' pbot' r }
 263
 264 getleft (UPortVar   v)  = v
 265 getleft (UPortPair a b) = getleft a
 266
 267 getright (UPortVar   v)  = v
 268 getright (UPortPair a b) = getright b
 269
 270 strip :: [(Ports,Int,Ports)] -> [(UVar,Int,UVar)]
 271 strip = map (\(v1,x,v2) -> (getright v1, x, getleft v2))
 272
 273
 274 -- must use bubblesort because the ordering isn't transitive
 275 sortit :: [(UVar,Int,UVar)] -> [(UVar,Int,UVar)]
 276 sortit x = let x' = stupidSort x [] in if x==x' then x else sortit x'
 277  where
 278   stupidSort []    x = x
 279   stupidSort (h:t) x = stupidSort t (stupidInsert h x)
 280
 281   stupidInsert t []    = [t]
 282   stupidInsert t@(_,_,t') ((a@(_,_,a')):b) = if cmp' x t' a' == LT
 283                                              then t:a:b
 284                                              else a:(stupidInsert t b)
 285
 286   cmp' :: [(UVar,Int,UVar)] -> UVar -> UVar -> Ordering
 287   cmp' [] a b = EQ -- compare a b
 288   cmp' ((v1,_,v2):r) a b = if a == v1 && b==v2
 289                            then LT
 290                            else if a == v2 && b==v1
 291                                 then GT
 292                                 else cmp' r a b
 293
 294 lookup'' :: Map UVar Int -> UVar -> Int
 295 lookup'' m x = case lookup x m of
 296                  Nothing -> 0
 297                  Just y -> y
 298
 299 valuatit :: Map UVar Int -> [(UVar,Int,UVar)] -> Map UVar Int
 300 valuatit m []            = m
 301 valuatit m ((v1,k,v2):r) = valuatit m' r
 302                             where
 303                               m'  = insert v2 v2' m
 304                               v2' = max (lookup'' m v2) (k+(lookup'' m v1))
 305
 306 resolve'k :: UyM () [(Ports,Int,Ports)]
 307 resolve'k = do { k  <- getK
 308               ; k' <- mapM (\(v1,x,v2) -> do { v1' <- getM v1 ; v2' <- getM v2 ; return (v1',x,v2') }) k
 309               ; return k'
 310               }
 311
 312 toTikZ :: GArrowSkeleton m a b -> String
 313 toTikZ g = tikZ m d 0
 314             where
 315               (_,_,_,(d,k)) = runU $ do { d <- ga2diag g
 316                                         ; d' <- resolve' d
 317                                         ; k' <- resolve'k
 318                                         ; return (d',k') }
 319               s = sortit (strip k)
 320               m = valuatit empty s
 321 toTikZ' :: GArrowSkeleton m a b -> String
 322 toTikZ' g = foldr (\x y -> x++"\\\\\n"++y) [] (map foo s)
 323              where
 324                (_,_,_,k) = runU $ ga2diag g >>= resolve' >>= \_ -> resolve'k
 325                foo (v1,v,v2) = "x_{"++show v1++"} + "++show v++" \\leq x_{"++show v2++"} = " ++ (show $ lookup'' m v2)
 326                s = sortit (strip k)
 327                m = valuatit empty s
 328
 329 tikz' :: (forall g a .
 330                  PGArrow g (GArrowUnit g) a ->
 331                  (
 332                    forall b . PGArrow g (GArrowTensor g b b) b) ->
 333                      PGArrow g (GArrowUnit g) a) -> IO ()
 334 tikz' x = tikz $ optimize $ unG (x (PGArrowD { unG = GAS_const 12 }) (PGArrowD { unG = GAS_merge }) )
 335 main = do putStrLn "hello"
 336 tikz example
 337      = do putStrLn "\\documentclass{article}"
 338           putStrLn "\\usepackage[paperwidth=\\maxdimen,paperheight=\\maxdimen]{geometry}"
 339           putStrLn "\\usepackage{tikz}"
 340           putStrLn "\\usepackage{amsmath}"
 341           putStrLn "\\usepackage[tightpage,active]{preview}"
 342           putStrLn "\\begin{document}"
 343           putStrLn "\\setlength\\PreviewBorder{5pt}"
 344           putStrLn "\\begin{preview}"
 345           putStrLn $ "\\begin{tikzpicture}[every on chain/.style={join=by ->},yscale=-1]"
 346           putStrLn (toTikZ example)
 347           putStrLn "\\end{tikzpicture}"
 348           putStrLn "\\end{preview}"
 349           --putStrLn "\\pagebreak"
 350           --putStrLn "\\begin{align*}"
 351           --putStr   (toTikZ' example)
 352           --putStrLn "\\end{align*}"
 353           putStrLn "\\end{document}"