src/Literate/Accounting/Chart.hs

   1 {-# LANGUAGE NoRebindableSyntax #-}
   2 {-# OPTIONS_GHC -Wno-orphans #-}
   3
   4 module Literate.Accounting.Chart where
   5
   6 import Control.Applicative (Applicative (..))
   7 import Control.DeepSeq (NFData (..))
   8 import Control.Monad (Monad (..))
   9 import Control.Monad.Trans.Writer qualified as MT
  10 import Data.Bool
  11 import Data.Eq (Eq (..))
  12 import Data.Foldable (Foldable (..), all)
  13 import Data.Function (const, flip, ($), (.))
  14 import Data.Functor (Functor (..), (<$>))
  15 import Data.List qualified as List
  16 import Data.List.NonEmpty (NonEmpty (..))
  17 import Data.List.NonEmpty qualified as NonEmpty
  18 import Data.Map.Strict (Map)
  19 import Data.Map.Strict qualified as Map
  20 import Data.Maybe (Maybe (..), isNothing)
  21 import Data.Monoid (Monoid (..))
  22 import Data.Ord (Ord (..))
  23 import Data.Semigroup (Semigroup (..))
  24 import Data.String (String)
  25 import Data.Traversable (Traversable (..))
  26 import Data.Tuple (fst, snd)
  27 import GHC.Exts qualified as GHC
  28 import Text.Show (Show (..))
  29
  30 import Literate.Accounting.Math
  31
  32 -- * Type 'Chart'
  33 newtype Chart k a = Chart {unChart :: Map.Map k (a, Chart k a)}
  34   deriving newtype (Eq, NFData)
  35 instance (Show k, Show a) => Show (Chart k a) where
  36   show = List.unlines . drawMap
  37    where
  38     drawNode :: (k, (a, Chart k a)) -> [String]
  39     drawNode (k, (a, ts0)) =
  40       List.zipWith (<>) (List.lines (show k)) (" " <> show a : List.repeat "")
  41         <> drawMap ts0
  42     drawMap = go . Map.toList . unChart
  43      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 =
  59     Chart $
  60       Map.unionWith
  61         (\new old -> (fst old <> fst new, snd old <> snd new))
  62         (unChart x)
  63         (unChart y)
  64 instance (Semigroup a, Ord k) => Monoid (Chart k a) where
  65   mempty = Chart Map.empty
  66 instance (Ord k, Addable a) => Addable (Chart k a) where
  67   x + y =
  68     Chart $
  69       Map.unionWith
  70         (\(ym, ya) (xm, xa) -> (xm + ym, xa + ya))
  71         (unChart x)
  72         (unChart y)
  73 instance (Ord k, Subable a) => Subable (Chart k a) where
  74   x - y =
  75     Chart $
  76       Map.unionWith
  77         (\(ym, ya) (xm, xa) -> (xm - ym, xa - ya))
  78         (unChart x)
  79         (unChart y)
  80
  81 -- ** Type 'ChartPath'
  82 type ChartPath = NonEmpty.NonEmpty
  83
  84 -- * Type 'Account'
  85 type Account = ChartPath
  86
  87 --newtype Account acct = Account (NonEmpty acct)
  88
  89 {- | @('insert' merge path value chart)@
  90 returns the given @chart@ with a mapping from @path@ to @value@,
  91 using @merge value oldValue@ in case @path@ already map to an @oldValue@.
  92 -}
  93 insert :: Ord k => a -> (a -> a -> a) -> ChartPath k -> a -> Chart k a -> Chart k a
  94 insert init merge p a ch = go ch p
  95  where
  96   go (Chart m) = \case
  97     k :| [] ->
  98       Chart $
  99         Map.insertWith
 100           (\_new (old, c) -> (merge a old, c))
 101           k
 102           (a, empty)
 103           m
 104     k :| k1 : ks ->
 105       Chart $
 106         Map.insertWith
 107           (\_new (old, c) -> (old, go c (k1 :| ks)))
 108           k
 109           (init, go empty (k1 :| ks))
 110           m
 111
 112 -- | Return the value (if any) associated with the given 'Path'.
 113 lookup :: Ord k => ChartPath k -> Chart k a -> Maybe a
 114 lookup (k :| ks) (Chart m) = do
 115   (a, ms) <- Map.lookup k m
 116   case ks of
 117     [] -> Just a
 118     (k' : ks') -> lookup (k' :| ks') ms
 119
 120 {- | @('filter' predicate chart)@ returns the given @chart@
 121 with only the paths whose mappings satisfied the given @predicate@.
 122 -}
 123 filter :: Ord k => (a -> Bool) -> Chart k a -> Chart k (Maybe a)
 124 filter f =
 125   Chart
 126     . Map.mapMaybe
 127       ( \(x, m) ->
 128           let fx = f x
 129            in let fm = filter f m
 130                in if not fx && all isNothing fm
 131                     then Nothing
 132                     else Just (if fx then Just x else Nothing, fm)
 133       )
 134     . unChart
 135
 136 -- | The empty 'Chart'.
 137 empty :: Chart k a
 138 empty = Chart Map.empty
 139
 140 {- | @('singleton' ancestorsValue path leafValue)@
 141 returns a 'Chart' mapping @path@ to @leafValue@
 142 and mapping all its ancestors paths to @ancestorsValue@.
 143 -}
 144 singleton :: Ord k => a -> ChartPath k -> a -> Chart k a
 145 singleton init ks a = insert init const ks a empty
 146
 147 {- | Return a 'Map' associating each 'ChartPath' in the given 'Chart',
 148 with its value mapped by the given function.
 149 -}
 150 flatten :: Ord k => (x -> y) -> Chart k x -> Map (ChartPath k) y
 151 flatten = flattenWithPath . const
 152
 153 flattenWithPath :: Ord k => ([k] -> x -> y) -> Chart k x -> Map (ChartPath k) y
 154 flattenWithPath = go []
 155  where
 156   go p f ch =
 157     Map.unions $
 158       Map.mapKeysMonotonic
 159         (NonEmpty.reverse . flip (:|) p)
 160         ( Map.mapWithKey
 161             ( \k (a, _children) ->
 162                 f (List.reverse (k : p)) a
 163             )
 164             (unChart ch)
 165         ) :
 166       Map.foldrWithKey
 167         (\k (_a, children) -> (go (k : p) f children :))
 168         []
 169         (unChart ch)
 170
 171 mapByDepthFirst :: Ord k => (Chart k b -> a -> b) -> Chart k a -> Chart k b
 172 mapByDepthFirst f =
 173   Chart
 174     . Map.map
 175       (\(a, ch) -> let m = mapByDepthFirst f ch in (f m a, m))
 176     . unChart
 177
 178 foldrPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> ChartPath k -> Chart k a -> acc -> acc
 179 foldrPath f = go [] . NonEmpty.toList
 180  where
 181   go _ [] _m acc = acc
 182   go p (k : ks) (Chart m) acc =
 183     case Map.lookup k m of
 184       Just (a, ch) -> f (NonEmpty.reverse (k :| p)) a $ go (k : p) ks ch acc
 185       Nothing -> acc
 186
 187 foldrWithPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> acc -> Chart k a -> acc
 188 foldrWithPath f = go []
 189  where
 190   go p acc =
 191     Map.foldrWithKey
 192       ( \k (a, ch) acc' ->
 193           f
 194             (NonEmpty.reverse (k :| p))
 195             a
 196             (go (k : p) acc' ch)
 197       )
 198       acc
 199       . unChart
 200 -- * Type 'ChartM'
 201
 202 -- | A 'Monad' to construct a 'Chart'.
 203 newtype ChartM k v m a = ChartM
 204   { unChartM ::
 205       MT.WriterT (v, Chart k v) m a
 206   }
 207   deriving newtype (Functor, Applicative, Monad)
 208
 209 runChartM :: Monad m => ChartM k v m a -> m (Chart k v)
 210 runChartM chM = do
 211   (_a, (_v, ch)) <- MT.runWriterT (unChartM chM)
 212   return ch
 213
 214 instance (Ord k, Monoid v, Monad m) => Semigroup (ChartM k v m a) where
 215   (<>) = (Control.Monad.>>)
 216 instance (Ord k, Monoid v, Monad m, Monoid a) => Monoid (ChartM k v m a) where
 217   mempty = return mempty
 218 instance
 219   ( Ord k
 220   , Monoid v
 221   , Monad m
 222   , Monoid a
 223   ) =>
 224   GHC.IsList (ChartM k v m a)
 225   where
 226   type Item (ChartM k v m a) = ChartM k v m a
 227   fromList = mconcat
 228   fromListN _n = mconcat
 229   toList = return