src/Haccounting/Chart.hs

   1 {-# LANGUAGE NoRebindableSyntax #-}
   2 {-# OPTIONS_GHC -Wno-orphans #-}
   3 module Haccounting.Chart where
   4
   5 import Control.Applicative (Applicative(..))
   6 import Control.DeepSeq (NFData(..))
   7 import Control.Monad (Monad(..))
   8 import Data.Bool
   9 import Data.Eq (Eq(..))
  10 import Data.Foldable (Foldable(..), all, foldr)
  11 import Data.Function (($), (.), const, flip)
  12 import Data.Functor (Functor(..), (<$>))
  13 import Data.List.NonEmpty (NonEmpty(..))
  14 import Data.Map.Strict (Map)
  15 import Data.Maybe (Maybe(..), isNothing)
  16 import Data.Monoid (Monoid(..))
  17 import Data.Ord (Ord(..))
  18 import Data.Ord (Ord(..))
  19 import Data.Semigroup (Semigroup(..))
  20 import Data.String (String)
  21 import Data.Traversable (Traversable(..))
  22 import Data.Tuple (fst, snd)
  23 import Prelude (error)
  24 import Text.Show (Show(..))
  25 import qualified Control.Monad.Trans.Writer as MT
  26 import qualified Data.List as List
  27 import qualified Data.List.NonEmpty as NonEmpty
  28 import qualified Data.Map.Strict as Map
  29 import qualified GHC.Exts as GHC
  30
  31 import Haccounting.Math
  32 import Haccounting.Rebindable
  33
  34 -- * Type 'Chart'
  35 newtype Chart k a = Chart { unChart :: Map.Map k (a, Chart k a) }
  36   deriving newtype (Eq, NFData)
  37 instance (Show k, Show a) => Show (Chart k a) where
  38   show = List.unlines . drawMap where
  39     drawNode :: (k, (a, Chart k a)) -> [String]
  40     drawNode (k, (a, ts0)) =
  41       List.zipWith (<>) (List.lines (show k)) (" " <> show a : List.repeat "") <>
  42       drawMap ts0
  43     drawMap = go . Map.toList . unChart where
  44       go [] = []
  45       go [t] = shift "` " "  " (drawNode t)
  46       go (t:ts) = shift "+ " "| " (drawNode t) <> go ts
  47     shift ind0 ind = List.zipWith (<>) (ind0 : List.repeat ind)
  48 instance Functor (Chart k) where
  49   fmap f = Chart . fmap (\(a, ch) -> (f a, fmap f ch)) . unChart
  50 instance Foldable (Chart k) where
  51   foldMap f = foldMap (\(a, ch) -> f a <> foldMap f ch) . unChart
  52 instance Traversable (Chart k) where
  53   traverse f =
  54     (Chart <$>) .
  55     traverse (\(a, ch) -> (,) <$> f a <*> traverse f ch) .
  56     unChart
  57 instance (Semigroup a, Ord k) => Semigroup (Chart k a) where
  58   x <> y = Chart $ Map.unionWith
  59     (\new old -> (fst old<>fst new, snd old<>snd new))
  60     (unChart x) (unChart y)
  61 instance (Semigroup a, Ord k) => Monoid (Chart k a) where
  62   mempty = Chart Map.empty
  63 instance (Ord k, Addable a) => Addable (Chart k a) where
  64   x + y = Chart $ Map.unionWith
  65     (\(ym, ya) (xm, xa) -> (xm + ym, xa + ya))
  66     (unChart x) (unChart y)
  67 instance (Ord k, Subable a) => Subable (Chart k a) where
  68   x - y = Chart $ Map.unionWith
  69     (\(ym, ya) (xm, xa) -> (xm - ym, xa - ya))
  70     (unChart x) (unChart y)
  71
  72 -- ** Type 'ChartPath'
  73 type ChartPath = NonEmpty.NonEmpty
  74
  75 insert :: Ord k => a -> (a -> a -> a) -> ChartPath k -> a -> Chart k a -> Chart k a
  76 insert init merge p a ch = go ch p
  77   where
  78   go (Chart m) = \case
  79     k:|[] -> Chart $ Map.insertWith
  80       (\_new (old, c) -> (merge a old, c))
  81       k (a, empty) m
  82     k:|k1:ks -> Chart $ Map.insertWith
  83       (\_new (old, c) -> (old, go c (k1:|ks)))
  84       k (init, go empty (k1:|ks)) m
  85
  86 -- | Return the value (if any) associated with the given 'Path'.
  87 lookup :: Ord k => ChartPath k -> Chart k a -> Maybe a
  88 lookup (k:|ks) (Chart m) = do
  89   (a, ms) <- Map.lookup k m
  90   case ks of
  91     [] -> Just a
  92     (k':ks') -> lookup (k':|ks') ms
  93
  94 filter :: Ord k => (a -> Bool) -> Chart k a -> Chart k (Maybe a)
  95 filter f =
  96   Chart . Map.mapMaybe (\(x, m) ->
  97     let fx = f x in
  98     let fm = filter f m in
  99     if not fx && all isNothing fm
 100     then Nothing
 101     else Just (if fx then Just x else Nothing, fm)
 102   ) . unChart
 103
 104 empty :: Chart k a
 105 empty = Chart Map.empty
 106
 107 singleton :: Ord k => a -> ChartPath k -> a -> Chart k a
 108 singleton init ks a = insert init const ks a empty
 109
 110 -- | Return a 'Map' associating each 'ChartPath' in the given 'Chart',
 111 -- with its value mapped by the given function.
 112 flatten :: Ord k => (x -> y) -> Chart k x -> Map (ChartPath k) y
 113 flatten = flattenWithPath . const
 114
 115 flattenWithPath :: Ord k => ([k] -> x -> y) -> Chart k x -> Map (ChartPath k) y
 116 flattenWithPath = go [] where
 117   go p f ch =
 118     Map.unions $
 119     Map.mapKeysMonotonic (NonEmpty.reverse . flip (:|) p) (
 120       Map.mapWithKey (\k (a, children) -> f (List.reverse (k : p)) a) (unChart ch)
 121     ) :
 122     Map.foldrWithKey
 123      (\k (_a, children) -> (go (k:p) f children :))
 124      [] (unChart ch)
 125
 126 mapByDepthFirst :: Ord k => (Chart k b -> a -> b) -> Chart k a -> Chart k b
 127 mapByDepthFirst f =
 128   Chart . Map.map
 129    (\(a, ch) -> let m = mapByDepthFirst f ch in (f m a, m)) .
 130   unChart
 131
 132 foldrPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> ChartPath k -> Chart k a -> acc -> acc
 133 foldrPath f = go [] . NonEmpty.toList where
 134   go _ [] _m acc = acc
 135   go p (k:ks) (Chart m) acc =
 136     case Map.lookup k m of
 137       Just (a, ch) -> f (NonEmpty.reverse (k:|p)) a $ go (k:p) ks ch acc
 138       Nothing -> acc
 139
 140 foldrWithPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> acc -> Chart k a -> acc
 141 foldrWithPath f = go [] where
 142   go p acc =
 143     Map.foldrWithKey (\k (a, ch) acc' ->
 144       f (NonEmpty.reverse (k:|p)) a
 145         (go (k:p) acc' ch)
 146       ) acc . unChart
 147
 148 -- * Type 'ChartM'
 149 -- | A 'Monad' to construct a 'Chart'.
 150 newtype ChartM k v m a = ChartM { unChartM ::
 151   MT.WriterT (v, Chart k v) m a
 152   } deriving newtype (Functor, Applicative, Monad)
 153
 154 runChartM :: Monad m => ChartM k v m a -> m (Chart k v)
 155 runChartM chM = do
 156   (_a, (_v, ch)) <- MT.runWriterT (unChartM chM)
 157   return ch
 158
 159 instance (Ord k, Monoid v, Monad m) => Semigroup (ChartM k v m a) where
 160   (<>) = (Control.Monad.>>)
 161 instance (Ord k, Monoid v, Monad m, Monoid a) => Monoid (ChartM k v m a) where
 162   mempty = return mempty
 163 instance
 164   ( Ord k
 165   , Monoid v
 166   , Monad m
 167   , Monoid a
 168   ) => GHC.IsList (ChartM k v m a) where
 169   type Item (ChartM k v m a) = ChartM k v m a
 170   fromList = mconcat
 171   fromListN _n = mconcat
 172   toList = return