Internal.hs

{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}

-- |
-- Module      :  Data.Tree.DUAL.Internal
-- Copyright   :  (c) 2011-2012 Brent Yorgey
-- License     :  BSD-style (see LICENSE)
-- Maintainer  :  diagrams-discuss@googlegroups.com
--
-- This module provides access to all of the internals of the
-- DUAL-tree implementation.  Depend on the internals at your own
-- risk!  For a safe public API (and complete documentation), see
-- "Data.Tree.DUAL".
--
-- The main things exported by this module which are not exported from
-- "Data.Tree.DUAL" are two extra types used in the implementation of
-- 'DUALTree', along with functions for manipulating them.  A type of
-- /non-empty/ trees, 'DUALTreeNE', is defined, as well as the type
-- 'DUALTreeU' which represents a non-empty tree paired with a cached
-- @u@ annotation.  'DUALTreeNE' and 'DUALTreeU' are mutually
-- recursive, so that recursive tree nodes are interleaved with cached
-- @u@ annotations.  'DUALTree' is defined by just wrapping
-- 'DUALTreeU' in 'Maybe.  This method has the advantage that the
-- type system enforces the invariant that there is only one
-- representation for the empty tree.  It also allows us to get away
-- with only 'Semigroup' constraints in many places.
module Data.Tree.DUAL.Internal (
  -- * DUAL-trees
  DUALTreeNE (..),
  DUALTreeU (..),
  DUALTree (..),

  -- * Constructing DUAL-trees
  empty,
  leaf,
  leafU,
  annot,
  applyD,

  -- * Modifying DUAL-trees
  applyUpre,
  applyUpost,
  mapUNE,
  mapUU,
  mapU,

  -- * Accessors and eliminators
  nonEmpty,
  getU,
  foldDUALNE,
  foldDUAL,
  flatten,
) where

import Control.Arrow ((***))
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NEL
import Data.Maybe (catMaybes, fromMaybe)
import Data.Monoid.Action
import Data.Semigroup
import Data.Typeable

------------------------------------------------------------
-- DUALTreeNE
------------------------------------------------------------

-- | /Non-empty/ DUAL-trees.
data DUALTreeNE d u a l
  = -- | Leaf with data value and @u@ annotation
    Leaf u l
  | -- | Leaf with only @u@ annotation
    LeafU u
  | -- | n-way branch, containing a /non-empty/ list
    --   of subtrees.
    Concat (NonEmpty (DUALTreeU d u a l))
  | -- | @d@ annotation
    Act d (DUALTreeU d u a l)
  | -- | Internal data value
    Annot a (DUALTreeU d u a l)
  deriving (Functor, Typeable, Show, Eq)

-- | The @Semigroup@ instance is carefully crafted so that 'sconcat'
--   creates a flat, rather than deeply nested, tree.
instance (Action d u, Semigroup u) => Semigroup (DUALTreeNE d u a l) where
  t1 <> t2 = sconcat (NEL.fromList [t1, t2])
  sconcat = Concat . NEL.map pullU

newtype DAct d = DAct {unDAct :: d}

instance
  (Semigroup d, Semigroup u, Action d u) =>
  Action (DAct d) (DUALTreeNE d u a l)
  where
  act (DAct d) (Act d' t) = Act (d <> d') t
  act (DAct d) t = Act d (pullU t)

------------------------------------------------------------
-- DUALTreeU
------------------------------------------------------------

-- | A non-empty DUAL-tree paired with a cached @u@ value.  These
--   should never be constructed directly; instead, use 'pullU'.
newtype DUALTreeU d u a l = DUALTreeU {unDUALTreeU :: (u, DUALTreeNE d u a l)}
  deriving (Functor, Typeable, Show, Eq)

-- | The @Semigroup@ instance is carefully crafted so that 'sconcat'
--   creates a flat, rather than deeply nested, tree.  Simply
--   auto-deriving the Semigroup instance produces an 'sconcat'
--   implementation which calls '(<>)' repeatedly, which is not what
--   we want.  Frustratingly, manually defining @sconcat =
--   DUALTreeU . sconcat . fmap unDUALTreeU@ does not work either,
--   since it ends up using 'sconcat' for pairs, which still turns
--   into nested calls to '(<>)'.
instance (Action d u, Semigroup u) => Semigroup (DUALTreeU d u a l) where
  DUALTreeU t1 <> DUALTreeU t2 = DUALTreeU (t1 <> t2)
  sconcat = DUALTreeU . (sconcat *** sconcat) . NEL.unzip . fmap unDUALTreeU

overDUALTreeU ::
  ((u, DUALTreeNE d u a l) -> (u', DUALTreeNE d u' a l)) ->
  DUALTreeU d u a l ->
  DUALTreeU d u' a l
overDUALTreeU f = DUALTreeU . f . unDUALTreeU

instance
  (Semigroup d, Semigroup u, Action d u) =>
  Action (DAct d) (DUALTreeU d u a l)
  where
  act d = overDUALTreeU (act (unDAct d) *** act d)

-- | \"Pull\" the root @u@ annotation out into a tuple.
pullU :: (Semigroup u, Action d u) => DUALTreeNE d u a l -> DUALTreeU d u a l
pullU t@(Leaf u _) = DUALTreeU (u, t)
pullU t@(LeafU u) = DUALTreeU (u, t)
pullU t@(Concat ts) = DUALTreeU (sconcat . NEL.map (fst . unDUALTreeU) $ ts, t)
pullU t@(Act d (DUALTreeU (u, _))) = DUALTreeU (act d u, t)
pullU t@(Annot _ (DUALTreeU (u, _))) = DUALTreeU (u, t)

------------------------------------------------------------
-- DUALTree
------------------------------------------------------------

-- | Rose (n-ary) trees with both upwards- (/i.e./ cached) and
--   downwards-traveling (/i.e./ accumulating) monoidal annotations.
--   Abstractly, a DUALTree is a rose (n-ary) tree with data (of type
--   @l@) at leaves, data (of type @a@) at internal nodes, and two
--   types of monoidal annotations, one (of type @u@) travelling
--   \"up\" the tree and one (of type @d@) traveling \"down\".  See
--   the documentation at the top of this file for full details.
--
--   @DUALTree@ comes with some instances:
--
--   * 'Functor', for modifying leaf data.  Note that 'fmap' of course
--     cannot alter any @u@ annotations.
--
--   * 'Semigroup'. @DUALTreeNE@s form a semigroup where @(\<\>)@
--     corresponds to adjoining two trees under a common parent root,
--     with @sconcat@ specialized to put all the trees under a single
--     parent.  Note that this does not satisfy associativity up to
--     structural equality, but only up to observational equivalence
--     under 'flatten'.  Technically using 'foldDUAL' directly enables
--     one to observe the difference, but it is understood that
--     'foldDUAL' should be used only in ways such that reassociation
--     of subtrees \"does not matter\".
--
--   * 'Monoid'. The identity is the empty tree.
newtype DUALTree d u a l = DUALTree {unDUALTree :: Maybe (DUALTreeU d u a l)}
  deriving (Functor, Typeable, Show, Eq)

-- | The @Semigroup@ instance is carefully crafted so that 'sconcat'
--   creates a flat, rather than deeply nested, tree.  Simply
--   auto-deriving the Semigroup instance produces an 'sconcat'
--   implementation which calls '(<>)' repeatedly, which is not what
--   we want.  Manually defining @sconcat = DUALTree . sconcat . fmap
--   unDUALTreeU@ does not even work either, since it ends up using
--   'sconcat' for 'Maybe', which still turns into nested calls to
--   '(<>)'.
instance (Action d u, Semigroup u) => Semigroup (DUALTree d u a l) where
  DUALTree t1 <> DUALTree t2 = DUALTree (t1 <> t2)

  sconcat = DUALTree . fmap sconcat . NEL.nonEmpty . catMaybes . NEL.toList . fmap unDUALTree

overDUALTree ::
  (Maybe (DUALTreeU d u a l) -> Maybe (DUALTreeU d u' a l)) ->
  DUALTree d u a l ->
  DUALTree d u' a l
overDUALTree f = DUALTree . f . unDUALTree

instance (Semigroup u, Action d u) => Monoid (DUALTree d u a l) where
  mempty = DUALTree mempty
  mappend = (<>)
  mconcat [] = mempty
  mconcat (x : xs) = sconcat (x :| xs)

-- | Apply a @d@ annotation at the root of a tree.  Semantically, all
--   @u@ annotations are transformed by the action of @d@, although
--   operationally @act@ incurs only a constant amount of work.
instance
  (Semigroup d, Semigroup u, Action d u) =>
  Action (DAct d) (DUALTree d u a l)
  where
  act = overDUALTree . fmap . act

------------------------------------------------------------
-- Convenience methods etc.
------------------------------------------------------------

-- | The empty DUAL-tree.  This is a synonym for 'mempty', but with a
--   more general type.
empty :: DUALTree d u a l
empty = DUALTree Nothing

-- | Construct a leaf node from a @u@ annotation along with a leaf
--   datum.
leaf :: u -> l -> DUALTree d u a l
leaf u l = DUALTree (Just (DUALTreeU (u, Leaf u l)))

-- | Construct a leaf node from a @u@ annotation.
leafU :: u -> DUALTree d u a l
leafU u = DUALTree (Just (DUALTreeU (u, LeafU u)))

-- | Add a @u@ annotation to the root, combining it (on the left) with
--   the existing cached @u@ annotation.  This function is provided
--   just for convenience; @applyUpre u t = 'leafU' u \<\> t@.
applyUpre :: (Semigroup u, Action d u) => u -> DUALTree d u a l -> DUALTree d u a l
applyUpre u t = leafU u <> t

-- | Add a @u@ annotation to the root, combining it (on the right) with
--   the existing cached @u@ annotation.  This function is provided
--   just for convenience; @applyUpost u t = t \<\> 'leafU' u@.
applyUpost :: (Semigroup u, Action d u) => u -> DUALTree d u a l -> DUALTree d u a l
applyUpost u t = t <> leafU u

-- | Add an internal data value at the root of a tree.  Note that this
--   only works on /non-empty/ trees; on empty trees this function is
--   the identity.
annot :: (Semigroup u, Action d u) => a -> DUALTree d u a l -> DUALTree d u a l
annot a = (overDUALTree . fmap) (pullU . Annot a)

-- | Apply a @d@ annotation at the root of a tree, transforming all
--   @u@ annotations by the action of @d@.
applyD ::
  (Semigroup d, Semigroup u, Action d u) =>
  d -> DUALTree d u a l -> DUALTree d u a l
applyD = act . DAct

-- | Decompose a DUAL-tree into either @Nothing@ (if empty) or a
--   top-level cached @u@ annotation paired with a non-empty
--   DUAL-tree.
nonEmpty :: DUALTree d u a l -> Maybe (u, DUALTreeNE d u a l)
nonEmpty = fmap unDUALTreeU . unDUALTree

-- | Get the @u@ annotation at the root, or @Nothing@ if the tree is
--   empty.
getU :: DUALTree d u a l -> Maybe u
getU = fmap fst . nonEmpty

------------------------------------------------------------
-- Maps
------------------------------------------------------------

-- XXX todo: try adding Map as a constructor, so we can delay the
-- mapping until the end too?

-- | Map a function (which must be a monoid homomorphism, and commute
--   with the action of @d@) over all the @u@ annotations in a non-empty
--   DUAL-tree.
mapUNE :: (u -> u') -> DUALTreeNE d u a l -> DUALTreeNE d u' a l
mapUNE f (Leaf u l) = Leaf (f u) l
mapUNE f (LeafU u) = LeafU (f u)
mapUNE f (Concat ts) = Concat ((NEL.map . mapUU) f ts)
mapUNE f (Act d t) = Act d (mapUU f t)
mapUNE f (Annot a t) = Annot a (mapUU f t)

-- | Map a function (which must be a monoid homomorphism, and commute
--   with the action of @d@) over all the @u@ annotations in a
--   non-empty DUAL-tree paired with its cached @u@ value.
mapUU :: (u -> u') -> DUALTreeU d u a l -> DUALTreeU d u' a l
mapUU f = overDUALTreeU (f *** mapUNE f)

-- | Map a function over all the @u@ annotations in a DUAL-tree.  The
--   function must be a monoid homomorphism, and must commute with the
--   action of @d@ on @u@.  That is, to use @mapU f@ safely it must be
--   the case that
--
--     * @f mempty == mempty@
--
--     * @f (u1 \<\> u2) == f u1 \<\> f u2@
--
--     * @f (act d u) == act d (f u)@
mapU :: (u -> u') -> DUALTree d u a l -> DUALTree d u' a l
mapU = overDUALTree . fmap . mapUU

------------------------------------------------------------
-- Folds
------------------------------------------------------------

-- | Fold for non-empty DUAL-trees.
foldDUALNE ::
  (Semigroup d, Monoid d) =>
  -- | Process a leaf datum along with the
  --   accumulation of @d@ values along the
  --   path from the root
  (d -> l -> r) ->
  -- | Replace @LeafU@ nodes
  r ->
  -- | Combine results at a branch node
  (NonEmpty r -> r) ->
  -- | Process an internal d node
  (d -> r -> r) ->
  -- | Process an internal datum
  (a -> r -> r) ->
  DUALTreeNE d u a l ->
  r
foldDUALNE = foldDUALNE' Nothing
 where
  foldDUALNE' dacc lf _ _ _ _ (Leaf _ l) = lf (fromMaybe mempty dacc) l
  foldDUALNE' _ _ lfU _ _ _ (LeafU _) = lfU
  foldDUALNE' dacc lf lfU con down ann (Concat ts) =
    con (NEL.map (foldDUALNE' dacc lf lfU con down ann . snd . unDUALTreeU) ts)
  foldDUALNE' dacc lf lfU con down ann (Act d t) =
    down d (foldDUALNE' (dacc <> Just d) lf lfU con down ann . snd . unDUALTreeU $ t)
  foldDUALNE' dacc lf lfU con down ann (Annot a t) =
    ann a (foldDUALNE' dacc lf lfU con down ann . snd . unDUALTreeU $ t)

-- | Fold for DUAL-trees. It is given access to the internal and leaf
--   data, internal @d@ values, and the accumulated @d@ values at each
--   leaf.  It is also allowed to replace \"@u@-only\" leaves with a
--   constant value.  In particular, however, it is /not/ given access
--   to any of the @u@ annotations, the idea being that those are used
--   only for /constructing/ trees.  If you do need access to @u@
--   values, you can duplicate the values you need in the internal
--   data nodes.
--
--   Be careful not to mix up the @d@ values at internal nodes with
--   the @d@ values at leaves.  Each @d@ value at a leaf satisfies the
--   property that it is the 'mconcat' of all internal @d@ values
--   along the path from the root to the leaf.
--
--   The result is @Nothing@ if and only if the tree is empty.
foldDUAL ::
  (Semigroup d, Monoid d) =>
  -- | Process a leaf datum along with the
  --   accumulation of @d@ values along the
  --   path from the root
  (d -> l -> r) ->
  -- | Replace @u@-only nodes
  r ->
  -- | Combine results at a branch node
  (NonEmpty r -> r) ->
  -- | Process an internal d node
  (d -> r -> r) ->
  -- | Process an internal datum
  (a -> r -> r) ->
  DUALTree d u a l ->
  Maybe r
foldDUAL _ _ _ _ _ (DUALTree Nothing) =
  Nothing
foldDUAL l u c d a (DUALTree (Just (DUALTreeU (_, t)))) =
  Just $ foldDUALNE l u c d a t

-- | A specialized fold provided for convenience: flatten a tree into
--   a list of leaves along with their @d@ annotations, ignoring
--   internal data values.
flatten :: (Semigroup d, Monoid d) => DUALTree d u a l -> [(l, d)]
flatten =
  fromMaybe []
    . foldDUAL
      (\d l -> [(l, d)])
      []
      (concat . NEL.toList)
      (\_ x -> x)
      (const id)