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