GC animation

GC-Anim.hs

{-# LANGUAGE ViewPatterns, GeneralizedNewtypeDeriving, DeriveDataTypeable#-}

import Diagrams.Prelude
import Diagrams.Backend.Cairo.CmdLine
import Diagrams.TwoD.Vector
import Diagrams.Core.Types
import qualified Data.Monoid.MList
import Data.Monoid.Coproduct

import Data.Typeable
import Data.Maybe
import Debug.Trace
import Control.Lens (set, ix, itraversed, indices, _1, _2, _3, _4, _5)
import Control.Arrow (first)

{-
getStyle :: Subdiagram b v m -> Style v
getStyle (Subdiagram _ (Option (Just a), _)) = killL a
-}

clWidth = 100
clHeight = 20
arrow_protrude = 25
arrow_extra = 6

closureField :: Int -> Maybe String -> Diagram B R2
closureField w mbs = (rect clWidth clHeight <> desc) # lw none # named (WordN w)
  where desc = case mbs of Nothing -> circle 3
                           Just s -> text s # fontSizeL 12

closure :: [Maybe String] -> Colour Double -> Diagram B R2
closure words c = translateY (-0.5*clHeight) $ lc c $
    r # stroke # lw thick # clipped r # alignB <>
    (vcat $ reverse $ zipWith closureField [0..] words) # alignB
  where
    w = clWidth
    h = fromIntegral (length words) * clHeight
    r = rect w h

cons :: Colour Double -> Diagram B R2
cons = closure [Just "Cons", Nothing, Nothing]

ind :: Colour Double -> Diagram B R2
ind = closure [Just "Ind", Nothing]

char :: Char -> Colour Double -> Diagram B R2
char c = closure [Just "Char", Just $ "'" ++ [c] ++ "'"] 

nil :: Colour Double -> Diagram B R2
nil = closure [Just "Nil"]

heap :: Diagram B R2
heap = vcat [phantom (nil grey) # named (Offset n) | n <- reverse [0..12]]

positionFor :: Int -> Int -> Diagram B R2 -> Diagram B R2
positionFor n c d = phantom d # named (positionAt n c)

positionAt :: Int -> Int -> Name
positionAt n c = toName (TP (n,c))

cellName :: Int -> Int -> Name
cellName n w = Closure n .> WordN w

realN :: Int -> Diagram B R2 -> Diagram B R2
realN n d = toName (Closure n) |> d

oldHeap :: Diagram B R2
oldHeap = Heap 1 |> heap

newHeap :: Diagram B R2
newHeap = Heap 2 |> heap

staticArea :: Diagram B R2
staticArea = Heap 3 |> heap

newtype TimePos = TP (Int, Int) deriving (Show, Ord, Eq, Typeable )
instance IsName TimePos

-- Names
newtype Heap = Heap Int deriving (Show, Ord, Eq, Typeable)
instance IsName Heap
newtype Offset = Offset Int deriving (Show, Ord, Eq, Typeable)
instance IsName Main.Offset
newtype Closure = Closure Int deriving (Show, Ord, Eq, Typeable)
instance IsName Closure
newtype WordN = WordN Int deriving (Show, Ord, Eq, Typeable)
instance IsName WordN

hPos :: Int -> Int -> Name
hPos n w = Heap n |> toName (Offset w)

data VH = V | H | P Double Double | PR Double Double -- horizontal first, verticla first, or following a path (with extra spacing and possibly reversed)
data Position
    = At Double Name
    | MoveFromTo Name Name VH Double
    | NotShown


animate :: Active a -> [a -> Active a] -> Active a
animate x [] = x
animate x (f:fs) = animate (x |>> f (activeEnd x)) fs

doWait :: a -> Active a
doWait x = x <$ ui

doMoveTo :: Name -> VH ->  Position -> Active Position
doMoveTo to vert p = MoveFromTo (endPos p) to vert <$> clamp ui

inBoth :: Applicative f => (a -> f a) -> (a, a) -> f (a, a)
inBoth f (p1,p2) = (,) <$> f p1 <*> f p2

endPos :: Position -> Name
endPos (At _ n) = n
endPos (MoveFromTo _ n _ _) = n
endPos NotShown = error "moveTo when not shown"

startAt :: Name -> Position
startAt  = At 1

-- This is a hack. We want something that takes now time, but changes endPos
doJumpTo :: Name -> Position -> Active Position
doJumpTo to _ = At 1 to <$ interval 0 0.01

fadeInAt :: Name -> Position -> Active Position
fadeInAt to _ = At <$> clamp ui <*> pure to

fadeIn :: Position -> Active Position
fadeIn p = At <$> clamp ui <*> pure (endPos p)

fadeOut :: Position -> Active Position
fadeOut p = At <$> backwards (clamp ui) <*> pure (endPos p)

type Arrow = (Double, Double, Colour Double, Name, Position)

type Scene = ([(Diagram B R2, Position)], [Arrow], Position)

startSpec :: Scene
startSpec =
          ( zipWith (first . realN) [1..] $
            [ (cons yellow,     startAt (hPos 1 10))
            , (char 'K' purple, startAt (hPos 1 8))
            , (cons green,      startAt (hPos 1 5))
            , (char 'F' blue,   startAt (hPos 1 3))
            , (cons brown,      startAt (hPos 1 0))
            , (ind grey,        NotShown)
            , (ind grey,        NotShown)
            , (ind grey,        NotShown)
            , (nil brown,       startAt (hPos 3 0))
            ]
          , [ (2, 0, yellow, cellName 1 1, startAt $ hPos 1 8)
            , (3, 0, yellow, cellName 1 2, startAt $ hPos 3 0)
            , (0, -1, green, cellName 3 1, startAt $ hPos 1 8)
            , (1, 0, green,  cellName 3 2, startAt $ hPos 1 10)
            , (0, 0, brown,  cellName 5 1, startAt $ hPos 1 3)
            , (1, 0, brown,  cellName 5 2, startAt $ hPos 1 5)
            , (4, 1, grey,   cellName 6 1, startAt $ cellName 3 0)
            , (5, 1, grey,   cellName 7 1, startAt $ cellName 2 0)
            , (6, 1, grey,   cellName 8 1, startAt $ cellName 1 0)
            , (0, -1, black, toName "live set",   startAt $ cellName 3 0)
            ]
            , startAt (toName "live set"))

positionSpec :: Active Scene
positionSpec = animate (doWait startSpec)
    [ _3              $ doMoveTo (cellName 3 0) V
    , _1 . ix 2 . _2  $ doMoveTo (hPos 2 0)     H
    , mapPair3A ( ix 5 . _2  $ fadeInAt (hPos 1 5)
                , ix 6 . _5  $ fadeIn
                , pure)
    , _3              $ doMoveTo (cellName 3 1) H
    , _3              $ doMoveTo (cellName 2 0) (P 0 (-1))
    , _2 . ix 2 . _5  $ doJumpTo (cellName 2 0)
    , _1 . ix 1 . _2  $ doMoveTo (hPos 2 3)     H
    , mapPair3A ( ix 6 . _2  $ fadeInAt (hPos 1 8)
                , ix 7 . _5  $ fadeIn
                , pure)
    , _3              $ doMoveTo (cellName 3 1) (PR 0 (-1))
    , _3              $ doMoveTo (cellName 3 2) V
    , _3              $ doMoveTo (cellName 1 0) (P 1 0)
    , _2 . ix 3 . _5  $ doJumpTo (cellName 1 0)
    , _1 . ix 0 . _2  $ doMoveTo (hPos 2 5)     H
    , mapPair3A ( ix 7 . _2  $ fadeInAt (hPos 1 10)
                , ix 8 . _5  $ fadeIn
                , pure)
    , _3              $ doMoveTo (cellName 1 1) H
    , _3              $ doMoveTo (cellName 7 0) (P 3 0)
    , (\f (a,b,c) -> (,,) <$> pure a <*> (ix 0 . _5) f b <*> f c) -- is this possible in a nicer way?
                      $ doMoveTo (cellName 7 1) V

    , (\f (a,b,c) -> (,,) <$> pure a <*> (ix 0 . _5) f b <*> f c)
                      $ doMoveTo (cellName 2 0) (P 5 1)
    , _3              $ doMoveTo (cellName 1 1) (PR 3 0)
    , _3              $ doMoveTo (cellName 1 2) V
    , doWait
    , doWait
    , doWait
    , mapPair3A ( itraversed . indices (`elem` [3,4,5,6,7]) . _2 $ fadeOut
                , itraversed . indices (`elem` [4,5,6,7,8]) . _5 $ fadeOut
                , fadeOut )
    ]

gcStar :: Diagram B R2
gcStar = star (StarSkip 2) (polygon (with & polyType .~ (PolyRegular 5 20))) # stroke # fc red # lw none


movements spec =
    applyAll $ reverse $
        [ d # moveToPos pos | (d, pos) <- objectSpec ] ++
        [ posToName pos $ \n2 ->
            myArrow ex exV (fc c . lc c . opacity (posOpacity pos)) n1 n2
        | (ex, exV, c, n1, pos) <- arrows] ++
        [ gcStar # moveToPos starSpec ]
  where
        (objectSpec, arrows, starSpec) = spec

posToName :: Position -> (Name -> Diagram B R2 -> Diagram B R2) -> (Diagram B R2 -> Diagram B R2)
posToName (At o i) cont = cont i
posToName NotShown cont = id
posToName (MoveFromTo i1 i2 _ r) cont | r <= 0 = cont i1
posToName (MoveFromTo i1 i2 _ r) cont | r >= 1 = cont i2
posToName (MoveFromTo i1 i2 hv r) cont =
    fromToTrail i1 i2 hv $ \trail ->
        let pos = trail `atParam` r
        in  cont (toName ()) . namePoint (const pos) ()


fromToTrail :: Name -> Name -> VH ->
    (Located (Trail R2) -> Diagram B R2 -> Diagram B R2) ->
    (Diagram B R2 -> Diagram B R2)
fromToTrail i1 i2 (P x xV)  cont = arrowTrail x xV i1 i2  cont
fromToTrail i1 i2 (PR x xV) cont = arrowTrail x xV i2 i1  (cont . reverseLocTrail)
fromToTrail i1 i2 vh cont =
    withName i1 $ \sub1 ->
    withName i2 $ \sub2 ->
        let pos1 = location sub1
            pos2 | V <- vh = pos1 |- pos3
                 | H <- vh = pos1 -| pos3
            pos3 = location sub2
            trail = fromVertices [pos1, pos2, pos3] :: Located (Trail R2)
        in (cont trail)

moveToPos :: Position -> Diagram B R2 -> (Diagram B R2 -> Diagram B R2)
moveToPos (At o i) d =
    withName i $ \sub -> atop $ d # opacity o # moveTo (location sub)
moveToPos (MoveFromTo i1 i2 vh r) d =
   fromToTrail i1 i2 vh $ \trail -> atop $ d # moveTo (trail `atParam` r)
moveToPos NotShown d = id

posOpacity :: Position -> Double
posOpacity (At o _) = o
posOpacity _  = 1



heaps :: Diagram B R2
heaps = hcat' (with & sep .~ clWidth)
    [ alignT $  strutY clHeight # named "live set"
                            ===
                (strutY clHeight <> text "live set" # fontSizeL 12)
    , alignB $ oldHeap # alignB === (strutY clHeight <> text "old heap" # fontSizeL 12)
    , alignB $ newHeap # alignB === (strutY clHeight <> text "new heap" # fontSizeL 12)
    , translateY (-2 * clHeight) $
      alignB $ staticArea # alignB === (strutY clHeight <> text "static area" # fontSizeL 12)
    ]

x ## f = f <$> x
infixl 8 ##

stableImage :: Animation B R2 -> Animation B R2
stableImage anim =
    anim ## withEnvelope (activeStart anim) ## frame 30 ## bg white

activeImage :: Animation B R2
activeImage = (movements <$> positionSpec <*> pure heaps) # stableImage ## font "Sans"

arrowTrail :: (IsName n1, IsName n2) => Double -> Double -> n1 -> n2 ->
    (Located (Trail R2) -> Diagram B R2 -> Diagram B R2) -> (Diagram B R2 -> Diagram B R2)
arrowTrail extra extraV n1 n2 cont =
  withName n1 $ \sub1 ->
  withName n2 $ \sub2 ->
    let s = location sub1
        e = location sub2 .+^ (extraV * arrow_extra) *^ unitY
        prod = (arrow_protrude + extra * arrow_extra) *^ unitX
        s2 = maybe s (.+^ prod) $ maxTraceP s unitX sub1

        fromRight = leftTurn unitY (e .-. s2)
        e3 | fromRight = fromMaybe e (maxTraceP e unitX sub2)
           | otherwise = fromMaybe e (maxTraceP e unit_X sub2)
        e2 = s2 |- e3
    in cont (fromVertices [s,s2,e2,e3])

myArrow :: (IsName n1, IsName n2) => Double -> Double -> (Diagram B R2 -> Diagram B R2) -> n1 -> n2 -> Diagram B R2 -> Diagram B R2
myArrow extra extraV style n1 n2 =
    arrowTrail extra extraV n1 n2 $ \trail -> atop $ mconcat
        [ circle 3 # moveTo (atStart trail) # stroke
        , trail # strokeLocTrail
        -- a bit hackish, normalAtEnd would be the clean solution
        , arrowBetween' (with & shaftStyle %~ lw none) (trail `atParam` 0.99) (atEnd trail)
        ] # style

(|-) :: P2 -> P2 -> P2
(unp2 -> (x,_)) |- (unp2 -> (_,y)) = p2 (x,y)

(-|) :: P2 -> P2 -> P2
(unp2 -> (_,y)) -| (unp2 -> (x,_)) = p2 (x,y)

mapPair3A :: Applicative f =>
    (a -> f a',  b -> f b', c -> f c') -> (a, b, c) -> f (a', b', c')
mapPair3A (f1, f2, f3) (a, b, c) = (,,) <$> f1 a <*> f2 b <*> f3 c

main :: IO ()
main = mainWith activeImage