7e24d4b36d94b0413d203e8ac511e6ae34d61df4
[darcs-mirror-arbtt.git] / src / LeftFold.hs
1 {-# LANGUAGE ExistentialQuantification, TypeOperators #-}
2
3 module LeftFold where
4
5 import Control.Applicative
6 import Data.List
7 import Data.Monoid
8 import Data.Strict ((:!:), Pair((:!:)))
9 import qualified Data.Strict as S
10 import qualified Data.Map.Strict as M
11 import Data.Maybe
12
13
14 data LeftFold x a = forall s. LeftFold {
15     start :: s,
16     process :: s -> x -> s,
17     finish :: s -> a
18     }
19     -- We keep things pure for as long as possible, to avoid constructing pairs
20     -- in <*> when not needed. Some of the more advanced code below (e.g.
21     -- intervals) is not properly tested with pure LeftFolds.
22     | Pure a 
23
24 leftFold :: a -> (a -> x -> a) -> LeftFold x a
25 leftFold s p = LeftFold s p id
26
27 instance Functor (LeftFold x) where
28     fmap f (Pure x) = Pure (f x)
29     fmap f (LeftFold st1 p1 f2) = LeftFold st1 p1 (f . f2)
30
31 instance Applicative (LeftFold x) where
32     pure x = Pure x
33     Pure f <*> c = f <$> c
34     LeftFold st1 p1 f1 <*> Pure x = LeftFold st1 p1 (\s -> f1 s x) 
35     LeftFold st1 p1 f1 <*> LeftFold st2 p2 f2 = LeftFold {
36         start   =                   st1 :!: st2,
37         process = \(s1 :!: s2) x -> p1 s1 x :!: p2 s2 x,
38         finish  = \(s1 :!: s2)   -> f1 s1 (f2 s2)
39         }
40
41 runLeftFold :: LeftFold x a -> [x] -> a
42 runLeftFold (Pure x) _ = x
43 runLeftFold (LeftFold st1 p1 f1) xs = f1 $! foldl' p1 st1 xs
44
45 monoidFold :: Monoid m => LeftFold m m
46 monoidFold = leftFold mempty mappend
47
48 mapElems :: LeftFold y a -> (x -> y) -> LeftFold x a 
49 mapElems (Pure x) _ = (Pure x)
50 mapElems (LeftFold s p f) t = LeftFold s (\s x -> p s $! t x) f
51
52 filterElems :: (x -> Bool) -> LeftFold x a -> LeftFold x a 
53 filterElems _ (Pure x) = (Pure x)
54 filterElems pred (LeftFold s p f) = LeftFold s (\s x -> if pred x then p s x else s) f
55
56 adjoin :: (x -> Bool) -> LeftFold (Bool :!: x) a -> LeftFold x a
57 adjoin p f = f `mapElems` (\x -> (p x :!: x))
58
59
60 onSelected :: LeftFold x a -> LeftFold (Bool :!: x) a
61 onSelected (Pure x) = Pure x
62 onSelected (LeftFold s p f) = LeftFold s (\s (b :!: x) -> if b then p s x else s) f
63
64 onJusts :: LeftFold x a -> LeftFold (Maybe x) a
65 onJusts (Pure x) = Pure x
66 onJusts (LeftFold s p f) = LeftFold s (\s mx -> maybe s (p s) mx) f
67
68 onAll :: LeftFold x a -> LeftFold (Bool :!: x) a
69 onAll (Pure x) = Pure x
70 onAll lf = lf `mapElems` S.snd
71
72 runOnGroups :: (x -> x -> Bool) -> LeftFold x y -> LeftFold y z -> LeftFold x z
73 runOnGroups eq _ (Pure ox) = Pure ox
74 runOnGroups eq (Pure ix) (LeftFold sto po fo) = LeftFold (S.Nothing :!: sto) go finish 
75     where go (S.Nothing :!: so) x             = (S.Just x :!: so)
76           go (S.Just x' :!: so) x | x `eq` x' = (S.Just x :!: so)
77                                   | otherwise = (S.Just x :!: po so ix)
78           finish (S.Nothing :!: so) = fo so
79           finish (S.Just _  :!: so) = fo (po so ix)
80 runOnGroups eq (LeftFold sti pi fi) (LeftFold sto po fo) = LeftFold (S.Nothing :!: sti :!: sto) go finish 
81     where go (S.Nothing :!: si :!: so) x             = (S.Just x :!: pi si x  :!: so)
82           go (S.Just x' :!: si :!: so) x | x `eq` x' = (S.Just x :!: pi si x  :!: so)
83                                          | otherwise = (S.Just x :!: pi sti x :!: po so (fi si))
84           finish (S.Nothing :!: si :!: so) = fo so
85           finish (S.Just _  :!: si :!: so) = fo (po so (fi si))
86
87 runOnIntervals :: LeftFold x y -> LeftFold y z -> LeftFold (Bool :!: x) z
88 runOnIntervals _ (Pure ox) = (Pure ox)
89 runOnIntervals (Pure ix) (LeftFold so po fo) = LeftFold (False :!: S.Nothing) go finish 
90     where go (True :!: so) (True :!: x)       = (True :!: so)
91           go (True :!: S.Just so) (False :!: x) = (False :!: S.Just (po so ix))
92           go (True :!: S.Nothing) (False :!: x) = (False :!: S.Just (po so ix))
93           go (False :!: so) (True :!: x)      = (True :!: so)
94           go (False :!: so) (False :!: x)     = (False :!: so)
95           finish (False :!: S.Just so) = fo so
96           finish (False :!: S.Nothing) = fo so
97           finish (True  :!: S.Just so) = fo (po so ix)
98           finish (True  :!: S.Nothing) = fo (po so ix)
99 runOnIntervals (LeftFold si pi fi) (LeftFold so po fo) = LeftFold (S.Nothing :!: S.Nothing) go finish 
100     where go (S.Just si :!: so) (True :!: x) = (S.Just (pi si x) :!: so)
101           go (S.Just si :!: S.Just so) (False :!: x) = (S.Nothing :!: S.Just (po so $! fi si))
102           go (S.Just si :!: S.Nothing) (False :!: x) = (S.Nothing :!: S.Just (po so $! fi si))
103           go (S.Nothing :!: so) (True :!: x) = (S.Just (pi si x) :!: so)
104           go (S.Nothing :!: so) (False :!: x) = (S.Nothing :!: so)
105           finish (S.Nothing :!: S.Just so) = fo so
106           finish (S.Nothing :!: S.Nothing) = fo so
107           finish (S.Just si :!: S.Just so) = fo (po so (fi si))
108           finish (S.Just si :!: S.Nothing) = fo (po so (fi si))
109
110 multiplex :: Ord k => (a -> k) -> LeftFold a b -> LeftFold a (M.Map k b)
111 multiplex key (LeftFold si pi fi) = LeftFold M.empty go finish
112     where go m x = M.alter go' (key x) m
113             where go' mbOld = Just $ pi (fromMaybe si mbOld) x
114           finish = M.map fi
115
116 lfLength :: LeftFold x Int
117 lfLength = leftFold 0 (\c _ -> c + 1)
118
119 lfFirst :: LeftFold x (Maybe x)
120 lfFirst = getFirst <$> monoidFold `mapElems` (First . Just)
121
122 lfLast :: LeftFold x (Maybe x)
123 lfLast = getLast <$> monoidFold `mapElems` (Last . Just)
124
125 toList :: LeftFold x [x]
126 toList = LeftFold [] (flip (:)) reverse
127
128 concatFold :: LeftFold [x] [x]
129 concatFold = concat <$> toList
130