src/Literate/Accounting/Chart.hs

   1 {-# LANGUAGE NoRebindableSyntax #-}
   2 {-# OPTIONS_GHC -Wno-orphans #-}
   3 module Literate.Accounting.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 Literate.Accounting.Math
  32 import Literate.Accounting.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 type Account = ChartPath
  76 -- * Type 'Account'
  77 --newtype Account acct = Account (NonEmpty acct)
  78
  79 insert :: Ord k => a -> (a -> a -> a) -> ChartPath k -> a -> Chart k a -> Chart k a
  80 insert init merge p a ch = go ch p
  81   where
  82   go (Chart m) = \case
  83     k:|[] -> Chart $ Map.insertWith
  84       (\_new (old, c) -> (merge a old, c))
  85       k (a, empty) m
  86     k:|k1:ks -> Chart $ Map.insertWith
  87       (\_new (old, c) -> (old, go c (k1:|ks)))
  88       k (init, go empty (k1:|ks)) m
  89
  90 -- | Return the value (if any) associated with the given 'Path'.
  91 lookup :: Ord k => ChartPath k -> Chart k a -> Maybe a
  92 lookup (k:|ks) (Chart m) = do
  93   (a, ms) <- Map.lookup k m
  94   case ks of
  95     [] -> Just a
  96     (k':ks') -> lookup (k':|ks') ms
  97
  98 filter :: Ord k => (a -> Bool) -> Chart k a -> Chart k (Maybe a)
  99 filter f =
 100   Chart . Map.mapMaybe (\(x, m) ->
 101     let fx = f x in
 102     let fm = filter f m in
 103     if not fx && all isNothing fm
 104     then Nothing
 105     else Just (if fx then Just x else Nothing, fm)
 106   ) . unChart
 107
 108 empty :: Chart k a
 109 empty = Chart Map.empty
 110
 111 singleton :: Ord k => a -> ChartPath k -> a -> Chart k a
 112 singleton init ks a = insert init const ks a empty
 113
 114 -- | Return a 'Map' associating each 'ChartPath' in the given 'Chart',
 115 -- with its value mapped by the given function.
 116 flatten :: Ord k => (x -> y) -> Chart k x -> Map (ChartPath k) y
 117 flatten = flattenWithPath . const
 118
 119 flattenWithPath :: Ord k => ([k] -> x -> y) -> Chart k x -> Map (ChartPath k) y
 120 flattenWithPath = go [] where
 121   go p f ch =
 122     Map.unions $
 123     Map.mapKeysMonotonic (NonEmpty.reverse . flip (:|) p) (
 124       Map.mapWithKey (\k (a, children) -> f (List.reverse (k : p)) a) (unChart ch)
 125     ) :
 126     Map.foldrWithKey
 127      (\k (_a, children) -> (go (k:p) f children :))
 128      [] (unChart ch)
 129
 130 mapByDepthFirst :: Ord k => (Chart k b -> a -> b) -> Chart k a -> Chart k b
 131 mapByDepthFirst f =
 132   Chart . Map.map
 133    (\(a, ch) -> let m = mapByDepthFirst f ch in (f m a, m)) .
 134   unChart
 135
 136 foldrPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> ChartPath k -> Chart k a -> acc -> acc
 137 foldrPath f = go [] . NonEmpty.toList where
 138   go _ [] _m acc = acc
 139   go p (k:ks) (Chart m) acc =
 140     case Map.lookup k m of
 141       Just (a, ch) -> f (NonEmpty.reverse (k:|p)) a $ go (k:p) ks ch acc
 142       Nothing -> acc
 143
 144 foldrWithPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> acc -> Chart k a -> acc
 145 foldrWithPath f = go [] where
 146   go p acc =
 147     Map.foldrWithKey (\k (a, ch) acc' ->
 148       f (NonEmpty.reverse (k:|p)) a
 149         (go (k:p) acc' ch)
 150       ) acc . unChart
 151
 152 -- * Type 'ChartM'
 153 -- | A 'Monad' to construct a 'Chart'.
 154 newtype ChartM k v m a = ChartM { unChartM ::
 155   MT.WriterT (v, Chart k v) m a
 156   } deriving newtype (Functor, Applicative, Monad)
 157
 158 runChartM :: Monad m => ChartM k v m a -> m (Chart k v)
 159 runChartM chM = do
 160   (_a, (_v, ch)) <- MT.runWriterT (unChartM chM)
 161   return ch
 162
 163 instance (Ord k, Monoid v, Monad m) => Semigroup (ChartM k v m a) where
 164   (<>) = (Control.Monad.>>)
 165 instance (Ord k, Monoid v, Monad m, Monoid a) => Monoid (ChartM k v m a) where
 166   mempty = return mempty
 167 instance
 168   ( Ord k
 169   , Monoid v
 170   , Monad m
 171   , Monoid a
 172   ) => GHC.IsList (ChartM k v m a) where
 173   type Item (ChartM k v m a) = ChartM k v m a
 174   fromList = mconcat
 175   fromListN _n = mconcat
 176   toList = return