/
DataToTree.hs
192 lines (155 loc) · 6.33 KB
/
DataToTree.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
{-# LANGUAGE ScopedTypeVariables #-}
module Test.SmartCheck.DataToTree
( forestReplaceChildren
, getAtIdx
, replaceAtIdx
, getIdxForest
, breadthLevels
, mkSubstForest
, depth
, tooDeep
) where
import Test.SmartCheck.Types
import Data.Tree
import Data.List
import Data.Maybe
import Data.Typeable
--------------------------------------------------------------------------------
-- Operations on Trees and Forests.
--------------------------------------------------------------------------------
-- | Return the list of values at each level in a Forest Not like levels in
-- Data.Tree (but what I imagined it should have done!).
breadthLevels :: Forest a -> [[a]]
breadthLevels forest =
takeWhile (not . null) go
where
go = map (getLevel forest) [0..]
--------------------------------------------------------------------------------
-- | Return the elements at level i from a forest. 0-based indexing.
getLevel :: Forest a -> Int -> [a]
getLevel fs 0 = map rootLabel fs
getLevel fs n = concatMap (\fs' -> getLevel (subForest fs') (n-1)) fs
--------------------------------------------------------------------------------
-- | Get the depth of a Forest. 0-based (an empty Forest has depth 0).
depth :: Forest a -> Int
depth forest = if null ls then 0 else maximum ls
where
ls = map depth' forest
depth' (Node _ []) = 1
depth' (Node _ forest') = 1 + depth forest'
--------------------------------------------------------------------------------
-- | How many members are at level i in the Tree?
levelLength :: Int -> Tree a -> Int
levelLength 0 t = length (subForest t)
levelLength n t = sum $ map (levelLength (n-1)) (subForest t)
--------------------------------------------------------------------------------
-- | Get the tree at idx in a forest. Nothing if the index is out-of-bounds.
getIdxForest :: Forest a -> Idx -> Maybe (Tree a)
getIdxForest forest (Idx (0 :: Int) n) =
if length forest > n then Just (forest !! n)
else Nothing
getIdxForest forest idx =
-- Should be a single Just x in the list, holding the value.
listToMaybe . catMaybes . snd $ acc
where
acc = mapAccumL findTree (column idx) (map Just forest)
l = level idx - 1
-- Invariant: not at the right level yet.
findTree :: Int -> Maybe (Tree a) -> (Int, Maybe (Tree a))
findTree n Nothing = (n, Nothing)
findTree n (Just t) =
let len = levelLength l t in
if n < 0 -- Already found index
then (n, Nothing)
else if n < len -- Big enough to index, so we climb down this one.
then let t' = getIdxForest (subForest t) (Idx l n) in
(n-len, t')
else (n-len, Nothing)
--------------------------------------------------------------------------------
-- Morally, we should be using generic zippers and a nice, recursive breadth-first search function, e.g.
{-
data Tree = N Int Tree Tree
| E
index :: Int -> Tree -> Tree
index = index' []
where
index' :: [Tree] -> Int -> Tree -> Tree
index' _ 0 t = t
index' [] idx (N i t0 t1) = index' [t1] (idx-1) t0
index' (k:ks) idx E = index' ks (idx-1) k
index' (k:ks) idx (N i t0 t1) = index' (ks ++ [t0, t1]) (idx-1) k
-}
-- | Returns the value at index idx. Returns nothing if the index is out of
-- bounds.
getAtIdx :: SubTypes a
=> a -- ^ Value
-> Idx -- ^ Index of hole
-> Maybe Int -- ^ Maximum depth we want to extract
-> Maybe SubT
getAtIdx d Idx { level = l, column = c } maxDepth
| tooDeep l maxDepth = Nothing
| length lev > c = Just (lev !! c)
| otherwise = Nothing
where
lev = getLevel (subTypes d) l
--------------------------------------------------------------------------------
tooDeep :: Int -> Maybe Int -> Bool
tooDeep l = maybe False (l >)
--------------------------------------------------------------------------------
data SubStrat = Parent -- ^ Replace everything in the path from the root to
-- here. Used as breadcrumbs to the value. Chop the
-- subforest.
| Children -- ^ Replace a value and all of its subchildren.
deriving (Show, Read, Eq)
--------------------------------------------------------------------------------
forestReplaceParent, forestReplaceChildren :: Forest a -> Idx -> a -> Forest a
forestReplaceParent = sub Parent
forestReplaceChildren = sub Children
--------------------------------------------------------------------------------
sub :: SubStrat -> Forest a -> Idx -> a -> Forest a
-- on right level, and we'll assume correct subtree.
sub strat forest (Idx (0 :: Int) n) a =
snd $ mapAccumL f 0 forest
where
f i node | i == n = ( i+1, news )
| otherwise = ( i+1, node )
where
news = case strat of
Parent -> Node a []
Children -> fmap (const a) (forest !! n)
sub strat forest idx a =
snd $ mapAccumL findTree (column idx) forest
where
l = level idx - 1
-- Invariant: not at the right level yet.
findTree n t
-- Already found index
| n < 0 = (n, t)
-- Big enough to index, so we climb down this one.
| n < len = (n-len, newTree)
| otherwise = (n-len, t)
where
len = levelLength l t
newTree = Node newRootLabel (sub strat (subForest t) (Idx l n) a)
newRootLabel = case strat of
Parent -> a
Children -> rootLabel t
--------------------------------------------------------------------------------
-- Operations on SubTypes.
--------------------------------------------------------------------------------
-- | Make a substitution Forest (all proper children). Initially we don't
-- replace anything.
mkSubstForest :: SubTypes a => a -> b -> Forest b
mkSubstForest a b = map tMap (subTypes a)
where tMap = fmap (const b)
--------------------------------------------------------------------------------
-- | Replace a value at index idx generically in a Tree/Forest generically.
replaceAtIdx :: (SubTypes a, Typeable b)
=> a -- ^ Parent value
-> Idx -- ^ Index of hole to replace
-> b -- ^ Value to replace with
-> Maybe a
replaceAtIdx m idx = replaceChild m (forestReplaceParent subF idx Subst)
where
subF = mkSubstForest m Keep
--------------------------------------------------------------------------------