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 import Prelude (error)
  30
  31 import GHC.Int (Int)
  32 import Literate.Accounting.Math
  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 = showChartAsTree
  39
  40 showChartAsTree :: Show k => Show a => Chart k a -> String
  41 showChartAsTree = List.unlines . drawMap
  42   where
  43     -- drawNode :: (k, (a, Chart k a)) -> [String]
  44     drawNode (k, (a, ts0)) =
  45       List.zipWith
  46         (<>)
  47         (List.lines (showsPrec 11 k ""))
  48         (List.lines (" " <> showsPrec 11 a "") <> List.repeat "")
  49         <> drawMap ts0
  50     drawMap = go . Map.toList . unChart
  51       where
  52         go [] = []
  53         go [t] = shift "` " "  " (drawNode t)
  54         go (t : ts) = shift "+ " "| " (drawNode t) <> go ts
  55     shift ind0 ind = List.zipWith (<>) (ind0 : List.repeat ind)
  56
  57 table :: [[String]] -> String
  58 table cells = List.unlines (List.foldr (\cell acc -> "|" <> cell <> acc) "|" <$> rows)
  59   where
  60     maxCols :: Int
  61     maxWidths :: [Int] -- for each row
  62     rows :: [[String]]
  63     (maxCols, maxWidths, rows) = List.foldr go (0, List.repeat 0, []) cells
  64
  65     go :: [String] -> (Int, [Int], [[String]]) -> (Int, [Int], [[String]])
  66     go row (accMaxCols, accMaxWidths, accRows) =
  67       ( max accMaxCols (List.length row)
  68       , List.zipWith max accMaxWidths (List.map List.length row <> List.repeat 0)
  69       , List.take maxCols (List.zipWith alignLeft row maxWidths) : accRows
  70       )
  71     alignRight cell maxWidth = List.replicate (maxWidth - List.length cell) ' ' <> cell
  72     alignLeft cell maxWidth = List.take maxWidth $ cell <> List.repeat ' '
  73
  74 {- >>> error $ table [["toto", "titi"], ["123", "4", "567890"], ["", "", "0"]]
  75 |toto|titi|
  76 |123 |4   |567890|
  77 |    |    |0     |
  78 -}
  79
  80 instance Functor (Chart k) where
  81   fmap f = Chart . fmap (\(a, ch) -> (f a, fmap f ch)) . unChart
  82 instance Foldable (Chart k) where
  83   foldMap f = foldMap (\(a, ch) -> f a <> foldMap f ch) . unChart
  84 instance Traversable (Chart k) where
  85   traverse f =
  86     (Chart <$>)
  87       . traverse (\(a, ch) -> (,) <$> f a <*> traverse f ch)
  88       . unChart
  89 instance (Semigroup a, Ord k) => Semigroup (Chart k a) where
  90   x <> y =
  91     Chart $
  92       Map.unionWith
  93         (\new old -> (fst old <> fst new, snd old <> snd new))
  94         (unChart x)
  95         (unChart y)
  96 instance (Semigroup a, Ord k) => Monoid (Chart k a) where
  97   mempty = Chart Map.empty
  98 instance (Ord k, Addable a) => Addable (Chart k a) where
  99   x + y =
 100     Chart $
 101       Map.unionWith
 102         (\(ym, ya) (xm, xa) -> (xm + ym, xa + ya))
 103         (unChart x)
 104         (unChart y)
 105 instance (Ord k, Subable a) => Subable (Chart k a) where
 106   x - y =
 107     Chart $
 108       Map.unionWith
 109         (\(ym, ya) (xm, xa) -> (xm - ym, xa - ya))
 110         (unChart x)
 111         (unChart y)
 112
 113 -- ** Type 'ChartPath'
 114 type ChartPath = NonEmpty.NonEmpty
 115
 116 -- * Type 'Account'
 117 type Account = ChartPath
 118
 119 --newtype Account acct = Account (NonEmpty acct)
 120
 121 {- | @('insert' merge path value chart)@
 122 returns the given @chart@ with a mapping from @path@ to @value@,
 123 using @merge value oldValue@ in case @path@ already map to an @oldValue@.
 124 -}
 125 insert :: Ord k => a -> (a -> a -> a) -> ChartPath k -> a -> Chart k a -> Chart k a
 126 insert init merge p a ch = go ch p
 127   where
 128     go (Chart m) = \case
 129       k :| [] ->
 130         Chart $
 131           Map.insertWith
 132             (\_new (old, c) -> (merge a old, c))
 133             k
 134             (a, empty)
 135             m
 136       k :| k1 : ks ->
 137         Chart $
 138           Map.insertWith
 139             (\_new (old, c) -> (old, go c (k1 :| ks)))
 140             k
 141             (init, go empty (k1 :| ks))
 142             m
 143
 144 -- | Return the value (if any) associated with the given 'Path'.
 145 lookup :: Ord k => ChartPath k -> Chart k a -> Maybe a
 146 lookup (k :| ks) (Chart m) = do
 147   (a, ms) <- Map.lookup k m
 148   case ks of
 149     [] -> Just a
 150     (k' : ks') -> lookup (k' :| ks') ms
 151
 152 {- | @('filter' predicate chart)@ returns the given @chart@
 153 with only the paths whose mappings satisfied the given @predicate@.
 154 -}
 155 filter :: Ord k => (a -> Bool) -> Chart k a -> Chart k (Maybe a)
 156 filter f =
 157   Chart
 158     . Map.mapMaybe
 159       ( \(x, m) ->
 160           let fx = f x
 161            in let fm = filter f m
 162                in if not fx && all isNothing fm
 163                     then Nothing
 164                     else Just (if fx then Just x else Nothing, fm)
 165       )
 166     . unChart
 167
 168 -- | The empty 'Chart'.
 169 empty :: Chart k a
 170 empty = Chart Map.empty
 171
 172 {- | @('singleton' ancestorsValue path leafValue)@
 173 returns a 'Chart' mapping @path@ to @leafValue@
 174 and mapping all its ancestors paths to @ancestorsValue@.
 175 -}
 176 singleton :: Ord k => a -> ChartPath k -> a -> Chart k a
 177 singleton init ks a = insert init const ks a empty
 178
 179 {- | Return a 'Map' associating each 'ChartPath' in the given 'Chart',
 180 with its value mapped by the given function.
 181 -}
 182 flatten :: Ord k => (x -> y) -> Chart k x -> Map (ChartPath k) y
 183 flatten = flattenWithPath . const
 184
 185 flattenWithPath :: Ord k => ([k] -> x -> y) -> Chart k x -> Map (ChartPath k) y
 186 flattenWithPath = go []
 187   where
 188     go p f ch =
 189       Map.unions $
 190         Map.mapKeysMonotonic
 191           (NonEmpty.reverse . flip (:|) p)
 192           ( Map.mapWithKey
 193               ( \k (a, _children) ->
 194                   f (List.reverse (k : p)) a
 195               )
 196               (unChart ch)
 197           ) :
 198         Map.foldrWithKey
 199           (\k (_a, children) -> (go (k : p) f children :))
 200           []
 201           (unChart ch)
 202
 203 mapByDepthFirst :: Ord k => (Chart k b -> a -> b) -> Chart k a -> Chart k b
 204 mapByDepthFirst f =
 205   Chart
 206     . Map.map
 207       (\(a, ch) -> let m = mapByDepthFirst f ch in (f m a, m))
 208     . unChart
 209
 210 foldrPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> ChartPath k -> Chart k a -> acc -> acc
 211 foldrPath f = go [] . NonEmpty.toList
 212   where
 213     go _ [] _m acc = acc
 214     go p (k : ks) (Chart m) acc =
 215       case Map.lookup k m of
 216         Just (a, ch) -> f (NonEmpty.reverse (k :| p)) a $ go (k : p) ks ch acc
 217         Nothing -> acc
 218
 219 foldrWithPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> acc -> Chart k a -> acc
 220 foldrWithPath f = go []
 221   where
 222     go p acc =
 223       Map.foldrWithKey
 224         ( \k (a, ch) acc' ->
 225             f
 226               (NonEmpty.reverse (k :| p))
 227               a
 228               (go (k : p) acc' ch)
 229         )
 230         acc
 231         . unChart
 232 -- * Type 'ChartM'
 233
 234 -- | A 'Monad' to construct a 'Chart'.
 235 newtype ChartM k v m a = ChartM
 236   { unChartM ::
 237       MT.WriterT (v, Chart k v) m a
 238   }
 239   deriving newtype (Functor, Applicative, Monad)
 240
 241 runChartM :: Monad m => ChartM k v m a -> m (Chart k v)
 242 runChartM chM = do
 243   (_a, (_v, ch)) <- MT.runWriterT (unChartM chM)
 244   return ch
 245
 246 instance (Ord k, Monoid v, Monad m) => Semigroup (ChartM k v m a) where
 247   (<>) = (Control.Monad.>>)
 248 instance (Ord k, Monoid v, Monad m, Monoid a) => Monoid (ChartM k v m a) where
 249   mempty = return mempty
 250 instance
 251   ( Ord k
 252   , Monoid v
 253   , Monad m
 254   , Monoid a
 255   ) =>
 256   GHC.IsList (ChartM k v m a)
 257   where
 258   type Item (ChartM k v m a) = ChartM k v m a
 259   fromList = mconcat
 260   fromListN _n = mconcat
 261   toList = return