{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies        #-}

{-# OPTIONS_GHC -fno-warn-unused-imports #-}
  -- for Data.Semigroup

-----------------------------------------------------------------------------
-- |
-- Module      :  Diagrams.Attributes.Compile
-- Copyright   :  (c) 2014 diagrams-lib team (see LICENSE)
-- License     :  BSD-style (see LICENSE)
-- Maintainer  :  diagrams-discuss@googlegroups.com
--
-- XXX
--
-----------------------------------------------------------------------------

module Diagrams.Attributes.Compile (
    SplitAttribute(..), splitAttr
  ) where

import           Data.Typeable

import           Control.Arrow       (second)
import           Control.Lens        ((%~), (&), _Wrapping')
import           Data.Kind           (Type)
import qualified Data.HashMap.Strict as HM
import           Data.Semigroup
import           Data.Tree           (Tree (..))

import           Diagrams.Core
import           Diagrams.Core.Style (Style (..), attributeToStyle)
import           Diagrams.Core.Types (RNode (..), RTree)

------------------------------------------------------------

-- This is a sort of roundabout, overly-general way to define
-- splitFills; it's done this way to facilitate testing.

class (AttributeClass (AttrType code), Typeable (PrimType code)) => SplitAttribute code where
  type AttrType code :: Type
  type PrimType code :: Type

  primOK :: code -> PrimType code -> Bool

-- | Push certain attributes down until they are at the roots of trees
--   containing only "safe" nodes.  In particular this is used to push
--   fill attributes down until they are over only loops; see
--   'splitFills'.
splitAttr :: forall code b v n a. SplitAttribute code => code -> RTree b v n a -> RTree b v n a
splitAttr :: code -> RTree b v n a -> RTree b v n a
splitAttr code
code = (RTree b v n a, Bool) -> RTree b v n a
forall a b. (a, b) -> a
fst ((RTree b v n a, Bool) -> RTree b v n a)
-> (RTree b v n a -> (RTree b v n a, Bool))
-> RTree b v n a
-> RTree b v n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe (AttrType code) -> RTree b v n a -> (RTree b v n a, Bool)
splitAttr' Maybe (AttrType code)
forall a. Maybe a
Nothing
  where

  -- splitAttr' is where the most interesting logic happens.
  -- Mutually recursive with splitAttr'Forest. rebuildNode and
  -- applyMfc are helper functions.
  --
  -- Input: attribute to apply to "safe" subtrees.
  --
  -- Output: tree with attributes pushed down appropriately, and
  -- a Bool indicating whether the tree contains only "safe" prims (True) or
  -- contains some unsafe ones (False).
  splitAttr' :: Maybe (AttrType code) -> RTree b v n a -> (RTree b v n a, Bool)

  -- RStyle node: Check for the special attribute, and split it out of
  -- the style, combining it with the incoming attribute.  Recurse and
  -- rebuild. The tricky bit is that we use some knot-tying to
  -- determine the right attribute to pass down to the subtrees based
  -- on this computed Bool: if all subtrees are safe, then we will
  -- apply the attribute at the root of this tree, and pass Nothing to
  -- all the subtrees.  Otherwise, we pass the given attribute along.
  -- This works out because the attribute does not need to be
  -- pattern-matched until actually applying it at some root, so the
  -- recursion can proceed and the Bool values be computed with the
  -- actual value of the attributes nodes filled in lazily.
  splitAttr' :: Maybe (AttrType code) -> RTree b v n a -> (RTree b v n a, Bool)
splitAttr' Maybe (AttrType code)
mattr (Node (RStyle Style v n
sty) Forest (RNode b v n a)
cs) = (RTree b v n a
t', Bool
ok)
    where
      mattr' :: Maybe (AttrType code)
mattr' = Maybe (AttrType code)
mattr Maybe (AttrType code)
-> Maybe (AttrType code) -> Maybe (AttrType code)
forall a. Semigroup a => a -> a -> a
<> Style v n -> Maybe (AttrType code)
forall a (v :: * -> *) n. AttributeClass a => Style v n -> Maybe a
getAttr Style v n
sty
      sty' :: Style v n
sty' = Style v n
sty Style v n -> (Style v n -> Style v n) -> Style v n
forall a b. a -> (a -> b) -> b
& (Unwrapped (Style v n) -> Style v n)
-> Iso' (Style v n) (Unwrapped (Style v n))
forall s. Wrapped s => (Unwrapped s -> s) -> Iso' s (Unwrapped s)
_Wrapping' Unwrapped (Style v n) -> Style v n
forall (v :: * -> *) n.
HashMap TypeRep (Attribute v n) -> Style v n
Style ((HashMap TypeRep (Attribute v n)
  -> Identity (HashMap TypeRep (Attribute v n)))
 -> Style v n -> Identity (Style v n))
-> (HashMap TypeRep (Attribute v n)
    -> HashMap TypeRep (Attribute v n))
-> Style v n
-> Style v n
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ TypeRep
-> HashMap TypeRep (Attribute v n)
-> HashMap TypeRep (Attribute v n)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
HM.delete TypeRep
ty
      ty :: TypeRep
ty   = AttrType code -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf (AttrType code
forall a. HasCallStack => a
undefined :: AttrType code)
      (Forest (RNode b v n a)
cs', Bool
ok) = Maybe (AttrType code)
-> Forest (RNode b v n a) -> (Forest (RNode b v n a), Bool)
splitAttr'Forest Maybe (AttrType code)
mattr' Forest (RNode b v n a)
cs
      t' :: RTree b v n a
t' | Bool
ok        = Maybe (AttrType code)
-> Bool -> RNode b v n a -> Forest (RNode b v n a) -> RTree b v n a
rebuildNode Maybe (AttrType code)
forall a. Maybe a
Nothing Bool
ok (Style v n -> RNode b v n a
forall b (v :: * -> *) n a. Style v n -> RNode b v n a
RStyle Style v n
sty) Forest (RNode b v n a)
cs'
         | Bool
otherwise = Maybe (AttrType code)
-> Bool -> RNode b v n a -> Forest (RNode b v n a) -> RTree b v n a
rebuildNode Maybe (AttrType code)
mattr Bool
ok (Style v n -> RNode b v n a
forall b (v :: * -> *) n a. Style v n -> RNode b v n a
RStyle Style v n
sty') Forest (RNode b v n a)
cs'

  -- RPrim node: check whether it
  --   * is some sort of prim not under consideration: don't apply the attribute; return True
  --   * is unsafe: don't apply the attribute; return False
  --   * is safe  :  do   apply the attribute; return True
  splitAttr' Maybe (AttrType code)
mattr (Node rp :: RNode b v n a
rp@(RPrim (Prim p
prm)) Forest (RNode b v n a)
_) =
      case p -> Maybe (PrimType code)
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast p
prm :: Maybe (PrimType code) of
        Maybe (PrimType code)
Nothing  -> (RNode b v n a -> Forest (RNode b v n a) -> RTree b v n a
forall a. a -> Forest a -> Tree a
Node RNode b v n a
rp [], Bool
True)
        Just PrimType code
p ->
          if code -> PrimType code -> Bool
forall code. SplitAttribute code => code -> PrimType code -> Bool
primOK code
code PrimType code
p
            then (Maybe (AttrType code)
-> Bool -> RNode b v n a -> Forest (RNode b v n a) -> RTree b v n a
rebuildNode Maybe (AttrType code)
mattr Bool
True RNode b v n a
rp [], Bool
True)
            else (RNode b v n a -> Forest (RNode b v n a) -> RTree b v n a
forall a. a -> Forest a -> Tree a
Node RNode b v n a
rp [], Bool
False)

  -- RFrozenTr, RAnnot, REmpty cases: just recurse and rebuild.  Note
  -- we assume that transformations do not affect the attributes.
  splitAttr' Maybe (AttrType code)
mattr (Node RNode b v n a
nd Forest (RNode b v n a)
cs)    = (RTree b v n a
t', Bool
ok)
    where
      (Forest (RNode b v n a)
cs', Bool
ok) = Maybe (AttrType code)
-> Forest (RNode b v n a) -> (Forest (RNode b v n a), Bool)
splitAttr'Forest Maybe (AttrType code)
mattr Forest (RNode b v n a)
cs
      t' :: RTree b v n a
t'        = Maybe (AttrType code)
-> Bool -> RNode b v n a -> Forest (RNode b v n a) -> RTree b v n a
rebuildNode Maybe (AttrType code)
mattr Bool
ok RNode b v n a
nd Forest (RNode b v n a)
cs'

  -- Recursively call splitAttr' on all subtrees, returning the
  -- logical AND of the Bool results returned (the whole forest is
  -- safe iff all subtrees are).
  splitAttr'Forest :: Maybe (AttrType code) -> [RTree b v n a] -> ([RTree b v n a], Bool)
  splitAttr'Forest :: Maybe (AttrType code)
-> Forest (RNode b v n a) -> (Forest (RNode b v n a), Bool)
splitAttr'Forest Maybe (AttrType code)
mattr Forest (RNode b v n a)
cs = (Forest (RNode b v n a)
cs', Bool
ok)
    where
      (Forest (RNode b v n a)
cs', Bool
ok) = ([Bool] -> Bool)
-> (Forest (RNode b v n a), [Bool])
-> (Forest (RNode b v n a), Bool)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Forest (RNode b v n a), [Bool])
 -> (Forest (RNode b v n a), Bool))
-> (Forest (RNode b v n a) -> (Forest (RNode b v n a), [Bool]))
-> Forest (RNode b v n a)
-> (Forest (RNode b v n a), Bool)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(RTree b v n a, Bool)] -> (Forest (RNode b v n a), [Bool])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(RTree b v n a, Bool)] -> (Forest (RNode b v n a), [Bool]))
-> (Forest (RNode b v n a) -> [(RTree b v n a, Bool)])
-> Forest (RNode b v n a)
-> (Forest (RNode b v n a), [Bool])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (RTree b v n a -> (RTree b v n a, Bool))
-> Forest (RNode b v n a) -> [(RTree b v n a, Bool)]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe (AttrType code) -> RTree b v n a -> (RTree b v n a, Bool)
splitAttr' Maybe (AttrType code)
mattr) (Forest (RNode b v n a) -> (Forest (RNode b v n a), Bool))
-> Forest (RNode b v n a) -> (Forest (RNode b v n a), Bool)
forall a b. (a -> b) -> a -> b
$ Forest (RNode b v n a)
cs

  -- Given a fill attribute, a Bool indicating whether the given
  -- subforest contains only loops, a node, and a subforest, rebuild a
  -- tree, applying the fill attribute as appropriate (only if the
  -- Bool is true and the attribute is not Nothing).
  rebuildNode :: Maybe (AttrType code) -> Bool -> RNode b v n a -> [RTree b v n a] -> RTree b v n a
  rebuildNode :: Maybe (AttrType code)
-> Bool -> RNode b v n a -> Forest (RNode b v n a) -> RTree b v n a
rebuildNode Maybe (AttrType code)
mattr Bool
ok RNode b v n a
nd Forest (RNode b v n a)
cs
    | Bool
ok        = Maybe (AttrType code) -> RTree b v n a -> RTree b v n a
applyMattr Maybe (AttrType code)
mattr (RNode b v n a -> Forest (RNode b v n a) -> RTree b v n a
forall a. a -> Forest a -> Tree a
Node RNode b v n a
nd Forest (RNode b v n a)
cs)
    | Bool
otherwise = RNode b v n a -> Forest (RNode b v n a) -> RTree b v n a
forall a. a -> Forest a -> Tree a
Node RNode b v n a
nd Forest (RNode b v n a)
cs

  -- Prepend a new fill color node if Just; the identity function if
  -- Nothing.
  applyMattr :: Maybe (AttrType code) -> RTree b v n a -> RTree b v n a
  applyMattr :: Maybe (AttrType code) -> RTree b v n a -> RTree b v n a
applyMattr Maybe (AttrType code)
Nothing  RTree b v n a
t = RTree b v n a
t
  applyMattr (Just AttrType code
a) RTree b v n a
t = RNode b v n a -> Forest (RNode b v n a) -> RTree b v n a
forall a. a -> Forest a -> Tree a
Node (Style v n -> RNode b v n a
forall b (v :: * -> *) n a. Style v n -> RNode b v n a
RStyle (Style v n -> RNode b v n a) -> Style v n -> RNode b v n a
forall a b. (a -> b) -> a -> b
$ Attribute v n -> Style v n
forall (v :: * -> *) n. Attribute v n -> Style v n
attributeToStyle (AttrType code -> Attribute v n
forall a (v :: * -> *) n. AttributeClass a => a -> Attribute v n
Attribute AttrType code
a)) [RTree b v n a
t]