lib/functional/iter_extras.nit

   1 # This file is part of NIT ( http://www.nitlanguage.org ).
   2 #
   3 # Copyright 2019 Louis-Vincent Boudreault <lv.boudreault95@gmail.com>
   4 #
   5 # Licensed under the Apache License, Version 2.0 (the "License");
   6 # you may not use this file except in compliance with the License.
   7 # You may obtain a copy of the License at
   8 #
   9 #     http://www.apache.org/licenses/LICENSE-2.0
  10 #
  11 # Unless required by applicable law or agreed to in writing, software
  12 # distributed under the License is distributed on an "AS IS" BASIS,
  13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14 # See the License for the specific language governing permissions and
  15 # limitations under the License.
  16
  17 # This modules provides a new functional interface for `Iterator`.
  18 module iter_extras
  19
  20 import functional_types
  21 import cartesian
  22
  23 redef class Array
  24
  25         # Sorts an array with a function.
  26         #
  27         # ~~~~nitish
  28         # class Person
  29         #       var name: String
  30         # end
  31         #
  32         # def get_name(p: Person) do return p.name
  33         #
  34         # var ps = [new Person("Turing"), new Person("Curry"), new Person("Alfredo")]
  35         # ps.sort_with(&get_name)
  36         # assert ps[0].name == "Alfredo"
  37         # assert ps[1].name == "Curry"
  38         # assert ps[2].name == "Turing"
  39         # ~~~~
  40         fun sort_with(f: Fun1[E, Comparable])
  41         do
  42                 var cmp = new ComparatorWith[E](f)
  43                 if length > 1 then
  44                         cmp.quick_sort(self, 0, length - 1)
  45                 end
  46         end
  47 end
  48
  49 redef interface Iterator[E]
  50
  51         # Applies a function to every elements
  52         #
  53         # ~~~~nitish
  54         # fun add(x: Int): Int do return x + 1
  55         #
  56         # var f = &add
  57         # var xs = [1,2,3,4,5]
  58         # var actual = xs.iterator.map(f).to_a
  59         # assert actual == [2,3,4,5,6]
  60         # ~~~~
  61         fun map(f: Fun1[E,Object]): MapIter[E,Object]
  62         do
  63                 return new MapIter[E,Object](self, f)
  64         end
  65
  66         # Iterator that gives the current count and element as a pair
  67         fun enumerate: EnumerateIter[E]
  68         do
  69                 return new EnumerateIter[E](self)
  70         end
  71
  72         # Iterator that filters elements by a predicate
  73         #
  74         # ~~~~nitish
  75         # fun lt10(x: Int): Bool do return x < 10
  76         #
  77         # var pred = &lt10
  78         # var xs = [1..20]
  79         # var actual = xs.iterator.filter(pred).to_a
  80         # assert actual == [1..9].to_a
  81         # ~~~~
  82         fun filter(pred: Fun1[E,Bool]): FilterIter[E]
  83         do
  84                 return new FilterIter[E](self,pred)
  85         end
  86
  87         # Checks if at least one element respects a predicate
  88         #
  89         # ~~~~nitish
  90         # fun eq10(x: Int): Bool do return x == 10
  91         #
  92         # var pred = &eq10
  93         # var xs = [1,2,5,7,9,10,44]
  94         # assert xs.iterator.any(pred)
  95         # var ys = []
  96         # assert not ys.iterator.any(pred)
  97         # assert not [1,2,44].iterator.any(pred)
  98         # ~~~~
  99         fun any(pred: Fun1[E,Bool]): Bool
 100         do
 101                 for x in self do
 102                         if pred.call(x) then
 103                                 return true
 104                         end
 105                 end
 106                 return false
 107         end
 108
 109         # Checks if all elements respect a predicate
 110         #
 111         # ~~~~nitish
 112         # fun lt10(x: Int): Bool do return x < 10
 113         #
 114         # var pred = &lt10
 115         # var xs = [1..9]
 116         # assert xs.iterator.all(pred)
 117         # assert [].iterator.all(pred)
 118         # assert not [1..10].iterator.all(pred)
 119         # ~~~~
 120         fun all(pred: Fun1[E,Bool]): Bool
 121         do
 122                 for x in self do
 123                         if not pred.call(x) then
 124                                 return false
 125                         end
 126                 end
 127                 return true
 128         end
 129
 130         # Folds an iterator from the left
 131         #
 132         # ~~~~nitish
 133         # fun adder(x: Int, y: Int): Int do return x + y
 134         #
 135         # var xs = [1..10]
 136         # assert xs.iterator.fold(0, &adder) == 55
 137         # ~~~~
 138         fun fold(acc: Object, f: Fun2[Object, E, Object]): Object
 139         do
 140                 for x in self do
 141                         acc = f.call(acc, x)
 142                 end
 143                 return acc
 144         end
 145
 146         # Folds and apply two element at a time
 147         #
 148         # ~~~~nitish
 149         # fun min_int(x: Int, y: Int): Int
 150         # do
 151         #       if x < y then return x
 152         #       return y
 153         # end
 154         #
 155         # var xs = [100,423,51,1,-19,55,999,-18]
 156         # assert xs.iterator.fold1(&min_int) == -19
 157         # ~~~~
 158         # REQUIRE : length > 1
 159         fun fold1(f: Fun2[E,E,E]): E
 160         do
 161                 var a1 = item
 162                 next
 163                 var a2 = item
 164                 next
 165                 var res = f.call(a1,a2)
 166                 for x in self do
 167                         res = f.call(res, x)
 168                 end
 169                 return res
 170         end
 171
 172         # Apply a mutation function over all elements
 173         #
 174         # ~~~~nitish
 175         # class Person
 176         #       var age: Int
 177         #       def incr_age
 178         #       do
 179         #               age += 1
 180         #       end
 181         # end
 182         #
 183         # var ps = [new Persone(1), new Person(2), new Person(3)]
 184         # var ages = ps.iterator.for_each(&Person::incr_age).map(&Person::age).to_a
 185         # assert ages == [2,3,4]
 186         # ~~~~
 187         fun for_each(f: Proc1[E])
 188         do
 189                 for x in self do
 190                         f.call(x)
 191                 end
 192         end
 193
 194         # Maps every element to a nested structure then flattens it
 195         #
 196         # ~~~~nitish
 197         # fun chars_fn(s: String): Iterator[Char]
 198         # do
 199         #       return s.chars.iterator
 200         # end
 201         # var cs = ["aaa","bbb","ccc"]
 202         # assert cs.iterator.flat_map(&chars_fn).to_a.join == "aaabbbccc"
 203         # ~~~~
 204         fun flat_map(f: Fun1[E, Iterator[Object]]): FlatMapIter[E, Object]
 205         do
 206                 return new FlatMapIter[E, Object](self, f)
 207         end
 208
 209         # Generates an `Iterator` whose elements are sorted by the function
 210         # passed in argument.
 211         #
 212         # ~~~~nitish
 213         # class Person
 214         #       var name: String
 215         # end
 216         #
 217         # def get_name(p: Person) do return p.name
 218         #
 219         # var ps = [new Person("Turing"), new Person("Curry"), new Person("Alfredo")]
 220         # var ordered_names = ps.iterator.order_by(&get_name).map(&get_name).to_a
 221         # assert ordered_names == ["Alfredo", "Curry", "Turing"]
 222         # ~~~~
 223         fun order_by(f: Fun1[E, Comparable]): OrderedIter[E]
 224         do
 225                 return new OrderedIter[E](self, f)
 226         end
 227
 228 end
 229
 230 # Base class for all iterators using functional types.
 231 private abstract class FunIter[OLD,NEW]
 232         super Iterator[NEW]
 233         var my_iter: Iterator[OLD]
 234
 235         redef fun next
 236         do
 237                 my_iter.next
 238         end
 239
 240         redef fun start
 241         do
 242                 my_iter.start
 243         end
 244
 245         redef fun finish
 246         do
 247                 my_iter.finish
 248         end
 249
 250         redef fun is_ok
 251         do
 252                 return my_iter.is_ok
 253         end
 254 end
 255
 256 # An iterator that maps each item with `f`.
 257 class MapIter[A,B]
 258         super FunIter[A,B]
 259         var f: Fun1[A, B]
 260
 261         redef fun item
 262         do
 263                 return f.call(my_iter.item)
 264         end
 265
 266 end
 267
 268 # An iterator that maps each item to a pair containing the item with its
 269 # current count.
 270 class EnumerateIter[E]
 271         super FunIter[E, Pair[Int,E]]
 272
 273         redef fun item
 274         do
 275                return new Pair[Int,E](0, my_iter.item)
 276         end
 277 end
 278
 279 # An tierator that filter its element by a predicate `pred`.
 280 class FilterIter[E]
 281         super FunIter[E,nullable E]
 282         var pred: Fun1[E, Bool]
 283
 284         redef init
 285         do
 286                 if is_ok and not pred.call(my_iter.item) then next
 287         end
 288
 289         redef fun item
 290         do
 291                 assert is_ok
 292                 return my_iter.item
 293         end
 294
 295         redef fun next
 296         do
 297                 loop
 298                         my_iter.next
 299                         if not is_ok then
 300                                 break
 301                         end
 302                         var x = my_iter.item
 303                         if pred.call(x) then
 304                                 break
 305                         end
 306                 end
 307         end
 308 end
 309
 310 # An iterator that maps each item to an iterator and yield
 311 # each item from it.
 312 class FlatMapIter[A,B]
 313         super FunIter[A,B]
 314         var f: Fun1[A, Iterator[B]]
 315         protected var inner: nullable Iterator[B] = null
 316
 317         redef init
 318         do
 319                 try_compute_inner
 320         end
 321
 322         redef fun item
 323         do
 324                 return inner.as(not null).item
 325         end
 326
 327         redef fun is_ok
 328         do
 329                 return inner != null
 330         end
 331
 332         redef fun next
 333         do
 334                 inner.next
 335                 if not inner.is_ok then
 336                         super
 337                         try_compute_inner
 338                 end
 339         end
 340
 341         # Tries to resolve the inner iterator.
 342         # Assigns null to `inner` if it fails.
 343         protected fun try_compute_inner
 344         do
 345                 inner = null
 346                 if not my_iter.is_ok then return
 347                 var res = f.call(my_iter.item)
 348                 if res.is_ok then inner = res
 349         end
 350 end
 351
 352 # An iterator that yield each item in order
 353 class OrderedIter[E]
 354         super FunIter[E,E]
 355         var f: Fun1[E, Comparable]
 356
 357         private var sorted_iter: Iterator[E] is noinit
 358         private var sorted_arr: Array[E] is noinit
 359
 360         redef init
 361         do
 362                 sorted_arr = my_iter.to_a
 363                 sorted_arr.sort_with(f)
 364                 sorted_iter = sorted_arr.iterator
 365         end
 366
 367         redef fun next
 368         do
 369                 sorted_iter.next
 370         end
 371
 372         redef fun item
 373         do
 374                 return sorted_iter.item
 375         end
 376
 377         redef fun is_ok
 378         do
 379                 return sorted_iter.is_ok
 380         end
 381
 382         redef fun finish
 383         do
 384                 sorted_iter.finish
 385         end
 386
 387         redef fun to_a
 388         do
 389                 return sorted_arr
 390         end
 391 end
 392
 393 # Comparator that use a function provided by the user to compare between elements.
 394 class ComparatorWith[E]
 395         super Comparator
 396         redef type COMPARED: E
 397
 398         var f: Fun1[E, Comparable]
 399
 400         redef fun compare(a,b)
 401         do
 402                 var x = f.call(a)
 403                 var y = f.call(b)
 404                 if x < y then return -1
 405                 if x > y then return 1
 406                 return 0
 407         end
 408 end