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, foldr)
  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 import Prelude (error)
  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
  39    where
  40     drawNode :: (k, (a, Chart k a)) -> [String]
  41     drawNode (k, (a, ts0)) =
  42       List.zipWith (<>) (List.lines (show k)) (" " <> show a : List.repeat "")
  43         <> drawMap ts0
  44     drawMap = go . Map.toList . unChart
  45      where
  46       go [] = []
  47       go [t] = shift "` " "  " (drawNode t)
  48       go (t : ts) = shift "+ " "| " (drawNode t) <> go ts
  49     shift ind0 ind = List.zipWith (<>) (ind0 : List.repeat ind)
  50 instance Functor (Chart k) where
  51   fmap f = Chart . fmap (\(a, ch) -> (f a, fmap f ch)) . unChart
  52 instance Foldable (Chart k) where
  53   foldMap f = foldMap (\(a, ch) -> f a <> foldMap f ch) . unChart
  54 instance Traversable (Chart k) where
  55   traverse f =
  56     (Chart <$>)
  57       . traverse (\(a, ch) -> (,) <$> f a <*> traverse f ch)
  58       . unChart
  59 instance (Semigroup a, Ord k) => Semigroup (Chart k a) where
  60   x <> y =
  61     Chart $
  62       Map.unionWith
  63         (\new old -> (fst old <> fst new, snd old <> snd new))
  64         (unChart x)
  65         (unChart y)
  66 instance (Semigroup a, Ord k) => Monoid (Chart k a) where
  67   mempty = Chart Map.empty
  68 instance (Ord k, Addable a) => Addable (Chart k a) where
  69   x + y =
  70     Chart $
  71       Map.unionWith
  72         (\(ym, ya) (xm, xa) -> (xm + ym, xa + ya))
  73         (unChart x)
  74         (unChart y)
  75 instance (Ord k, Subable a) => Subable (Chart k a) where
  76   x - y =
  77     Chart $
  78       Map.unionWith
  79         (\(ym, ya) (xm, xa) -> (xm - ym, xa - ya))
  80         (unChart x)
  81         (unChart y)
  82
  83 -- ** Type 'ChartPath'
  84 type ChartPath = NonEmpty.NonEmpty
  85
  86 type Account = ChartPath
  87 -- * Type 'Account'
  88
  89 --newtype Account acct = Account (NonEmpty acct)
  90
  91 insert :: Ord k => a -> (a -> a -> a) -> ChartPath k -> a -> Chart k a -> Chart k a
  92 insert init merge p a ch = go ch p
  93  where
  94   go (Chart m) = \case
  95     k :| [] ->
  96       Chart $
  97         Map.insertWith
  98           (\_new (old, c) -> (merge a old, c))
  99           k
 100           (a, empty)
 101           m
 102     k :| k1 : ks ->
 103       Chart $
 104         Map.insertWith
 105           (\_new (old, c) -> (old, go c (k1 :| ks)))
 106           k
 107           (init, go empty (k1 :| ks))
 108           m
 109
 110 -- | Return the value (if any) associated with the given 'Path'.
 111 lookup :: Ord k => ChartPath k -> Chart k a -> Maybe a
 112 lookup (k :| ks) (Chart m) = do
 113   (a, ms) <- Map.lookup k m
 114   case ks of
 115     [] -> Just a
 116     (k' : ks') -> lookup (k' :| ks') ms
 117
 118 filter :: Ord k => (a -> Bool) -> Chart k a -> Chart k (Maybe a)
 119 filter f =
 120   Chart
 121     . Map.mapMaybe
 122       ( \(x, m) ->
 123           let fx = f x
 124            in let fm = filter f m
 125                in if not fx && all isNothing fm
 126                     then Nothing
 127                     else Just (if fx then Just x else Nothing, fm)
 128       )
 129     . unChart
 130
 131 empty :: Chart k a
 132 empty = Chart Map.empty
 133
 134 singleton :: Ord k => a -> ChartPath k -> a -> Chart k a
 135 singleton init ks a = insert init const ks a empty
 136
 137 {- | Return a 'Map' associating each 'ChartPath' in the given 'Chart',
 138  with its value mapped by the given function.
 139 -}
 140 flatten :: Ord k => (x -> y) -> Chart k x -> Map (ChartPath k) y
 141 flatten = flattenWithPath . const
 142
 143 flattenWithPath :: Ord k => ([k] -> x -> y) -> Chart k x -> Map (ChartPath k) y
 144 flattenWithPath = go []
 145  where
 146   go p f ch =
 147     Map.unions $
 148       Map.mapKeysMonotonic
 149         (NonEmpty.reverse . flip (:|) p)
 150         ( Map.mapWithKey (\k (a, children) -> f (List.reverse (k : p)) a) (unChart ch)
 151         ) :
 152       Map.foldrWithKey
 153         (\k (_a, children) -> (go (k : p) f children :))
 154         []
 155         (unChart ch)
 156
 157 mapByDepthFirst :: Ord k => (Chart k b -> a -> b) -> Chart k a -> Chart k b
 158 mapByDepthFirst f =
 159   Chart
 160     . Map.map
 161       (\(a, ch) -> let m = mapByDepthFirst f ch in (f m a, m))
 162     . unChart
 163
 164 foldrPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> ChartPath k -> Chart k a -> acc -> acc
 165 foldrPath f = go [] . NonEmpty.toList
 166  where
 167   go _ [] _m acc = acc
 168   go p (k : ks) (Chart m) acc =
 169     case Map.lookup k m of
 170       Just (a, ch) -> f (NonEmpty.reverse (k :| p)) a $ go (k : p) ks ch acc
 171       Nothing -> acc
 172
 173 foldrWithPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> acc -> Chart k a -> acc
 174 foldrWithPath f = go []
 175  where
 176   go p acc =
 177     Map.foldrWithKey
 178       ( \k (a, ch) acc' ->
 179           f
 180             (NonEmpty.reverse (k :| p))
 181             a
 182             (go (k : p) acc' ch)
 183       )
 184       acc
 185       . unChart
 186 -- * Type 'ChartM'
 187
 188 -- | A 'Monad' to construct a 'Chart'.
 189 newtype ChartM k v m a = ChartM
 190   { unChartM ::
 191       MT.WriterT (v, Chart k v) m a
 192   }
 193   deriving newtype (Functor, Applicative, Monad)
 194
 195 runChartM :: Monad m => ChartM k v m a -> m (Chart k v)
 196 runChartM chM = do
 197   (_a, (_v, ch)) <- MT.runWriterT (unChartM chM)
 198   return ch
 199
 200 instance (Ord k, Monoid v, Monad m) => Semigroup (ChartM k v m a) where
 201   (<>) = (Control.Monad.>>)
 202 instance (Ord k, Monoid v, Monad m, Monoid a) => Monoid (ChartM k v m a) where
 203   mempty = return mempty
 204 instance
 205   ( Ord k
 206   , Monoid v
 207   , Monad m
 208   , Monoid a
 209   ) =>
 210   GHC.IsList (ChartM k v m a)
 211   where
 212   type Item (ChartM k v m a) = ChartM k v m a
 213   fromList = mconcat
 214   fromListN _n = mconcat
 215   toList = return