src/Symantic/Compta/Calc/Chart.hs

   1 {-# LANGUAGE NoRebindableSyntax #-}
   2 {-# OPTIONS_GHC -Wno-orphans #-}
   3 module Symantic.Compta.Calc.Chart where
   4
   5 import Control.Applicative (Applicative(..))
   6 import Control.DeepSeq (NFData(..))
   7 import Data.Bool
   8 import Data.Traversable (Traversable(..))
   9 import Data.Eq (Eq(..))
  10 import Data.Function (($), (.), const, flip)
  11 import Data.Functor (Functor(..), (<$>))
  12 import Data.Foldable (Foldable(..), all)
  13 import Data.Map.Strict (Map)
  14 import Data.Ord (Ord(..))
  15 import Data.Semigroup (Semigroup(..))
  16 import Data.Maybe (Maybe(..), isNothing)
  17 import Data.Monoid (Monoid(..))
  18 import Text.Show (Show(..))
  19 import Data.String (String)
  20 import Data.Tuple (fst, snd)
  21 import qualified Data.List as List
  22 import qualified Data.List.NonEmpty as NonEmpty
  23 import Data.List.NonEmpty (NonEmpty(..))
  24 import qualified Data.Map.Strict as Map
  25
  26 --import Symantic.Compta.Lang.Rebindable
  27 import Symantic.Compta.Lang.Math
  28
  29 -- * Type 'Chart'
  30 newtype Chart k a = Chart { unChart :: Map.Map k (a, Chart k a) }
  31   deriving newtype (Eq, NFData)
  32 instance (Show k, Show a) => Show (Chart k a) where
  33   show = List.unlines . drawMap where
  34     drawNode :: (k, (a, Chart k a)) -> [String]
  35     drawNode (k, (a, ts0)) =
  36       List.zipWith (<>) (List.lines (show k)) (" " <> show a : List.repeat "") <>
  37       drawMap ts0
  38     drawMap = go . Map.toList . unChart where
  39       go [] = []
  40       go [t] = shift "` " "  " (drawNode t)
  41       go (t:ts) = shift "+ " "| " (drawNode t) <> go ts
  42     shift ind0 ind = List.zipWith (<>) (ind0 : List.repeat ind)
  43 instance Functor (Chart k) where
  44   fmap f = Chart . fmap (\(a, ch) -> (f a, fmap f ch)) . unChart
  45 instance Foldable (Chart k) where
  46   foldMap f = foldMap (\(a, ch) -> f a <> foldMap f ch) . unChart
  47 instance Traversable (Chart k) where
  48   traverse f =
  49     (Chart <$>) .
  50     traverse (\(a, ch) -> (,) <$> f a <*> traverse f ch) .
  51     unChart
  52 instance (Semigroup a, Ord k) => Semigroup (Chart k a) where
  53   x <> y = Chart $ Map.unionWith
  54     (\new old -> (fst old<>fst new, snd old<>snd new))
  55     (unChart x) (unChart y)
  56 instance (Semigroup a, Ord k) => Monoid (Chart k a) where
  57   mempty = Chart Map.empty
  58 instance (Ord k, Addable a) => Addable (Chart k a) where
  59   x + y = Chart $ Map.unionWith
  60     (\(ym, ya) (xm, xa) -> (xm + ym, xa + ya))
  61     (unChart x) (unChart y)
  62 instance (Ord k, Subable a) => Subable (Chart k a) where
  63   x - y = Chart $ Map.unionWith
  64     (\(ym, ya) (xm, xa) -> (xm - ym, xa - ya))
  65     (unChart x) (unChart y)
  66
  67 -- ** Type 'ChartPath'
  68 type ChartPath = NonEmpty.NonEmpty
  69
  70 insert :: Ord k => a -> (a -> a -> a) -> ChartPath k -> a -> Chart k a -> Chart k a
  71 insert init merge p a ch = go ch p
  72   where
  73   go (Chart m) = \case
  74     k:|[] -> Chart $ Map.insertWith
  75       (\_new (old, c) -> (merge a old, c))
  76       k (a, empty) m
  77     k:|k1:ks -> Chart $ Map.insertWith
  78       (\_new (old, c) -> (old, go c (k1:|ks)))
  79       k (init, go empty (k1:|ks)) m
  80
  81 -- | Return the value (if any) associated with the given 'Path'.
  82 lookup :: Ord k => ChartPath k -> Chart k a -> Maybe a
  83 lookup (k:|ks) (Chart m) = do
  84   (a, ms) <- Map.lookup k m
  85   case ks of
  86     [] -> Just a
  87     (k':ks') -> lookup (k':|ks') ms
  88
  89 filter :: Ord k => (a -> Bool) -> Chart k a -> Chart k (Maybe a)
  90 filter f =
  91   Chart . Map.mapMaybe (\(x, m) ->
  92     let fx = f x in
  93     let fm = filter f m in
  94     if not fx && all isNothing fm
  95     then Nothing
  96     else Just (if fx then Just x else Nothing, fm)
  97   ) . unChart
  98
  99 empty :: Chart k a
 100 empty = Chart Map.empty
 101
 102 singleton :: Ord k => a -> ChartPath k -> a -> Chart k a
 103 singleton init ks a = insert init const ks a empty
 104
 105 -- | Return a 'Map' associating each 'ChartPath' in the given 'Chart',
 106 -- with its value mapped by the given function.
 107 flatten :: Ord k => (x -> y) -> Chart k x -> Map (ChartPath k) y
 108 flatten = flattenWithPath . const
 109
 110 flattenWithPath :: Ord k => ([k] -> x -> y) -> Chart k x -> Map (ChartPath k) y
 111 flattenWithPath = go [] where
 112   go p f ch =
 113     Map.unions $
 114     Map.mapKeysMonotonic (NonEmpty.reverse . flip (:|) p) (
 115       Map.mapWithKey (\k (a, children) -> f (List.reverse (k : p)) a) (unChart ch)
 116     ) :
 117     Map.foldrWithKey
 118      (\k (_a, children) -> (go (k:p) f children :))
 119      [] (unChart ch)
 120
 121 mapByDepthFirst :: Ord k => (Chart k b -> a -> b) -> Chart k a -> Chart k b
 122 mapByDepthFirst f =
 123   Chart . Map.map
 124    (\(a, ch) -> let m = mapByDepthFirst f ch in (f m a, m)) .
 125   unChart
 126
 127 foldrPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> ChartPath k -> Chart k a -> acc -> acc
 128 foldrPath f = go [] . NonEmpty.toList where
 129   go _ [] _m acc = acc
 130   go p (k:ks) (Chart m) acc =
 131     case Map.lookup k m of
 132       Just (a, ch) -> f (NonEmpty.reverse (k:|p)) a $ go (k:p) ks ch acc
 133       Nothing -> acc
 134
 135 foldrWithPath :: Ord k => (ChartPath k -> a -> acc -> acc) -> acc -> Chart k a -> acc
 136 foldrWithPath f = go [] where
 137   go p acc =
 138     Map.foldrWithKey (\k (a, ch) acc' ->
 139       f (NonEmpty.reverse (k:|p)) a
 140         (go (k:p) acc' ch)
 141       ) acc . unChart