more flexible CCList.sublists_of_len

src/core/CCList.ml

@@ -528,7 +528,7 @@ let take_drop n l = take n l, drop n l
     l1 @ l2 = l )
 *)
 
-let sublists_of_len ?(drop_short=true) ?offset n l =
+let sublists_of_len ?(last=fun _ -> None) ?offset n l =
   if n < 1 then invalid_arg "sublists_of_len: n must be > 0";
   let offset = match offset with
     | None -> n
@@ -538,11 +538,12 @@ let sublists_of_len ?(drop_short=true) ?offset n l =
   (* add sub-lists of [l] to [acc] *)
   let rec aux acc l =
     let group = take n l in
-    if List.length group < n
-    then (
-      (* this was the last group *)
-      if drop_short || group=[] then acc else group :: acc
-    ) else (
+    if group=[] then acc (* this was the last group, we are done *)
+    else if List.length group < n (* last group, with missing elements *)
+    then match last group with
+      | None -> acc
+      | Some group' -> group' :: acc
+    else (
       let l' = drop offset l in
       aux (group :: acc) l' (* continue *)
     )
@@ -556,7 +557,7 @@ let sublists_of_len ?(drop_short=true) ?offset n l =
   [[1;2];[3;4]] (subs 2 ~offset:2 [1;2;3;4])
   [[1;2];[2;3]] (subs 2 ~offset:1 [1;2;3])
   [[1;2];[4;5]] (subs 2 ~offset:3 [1;2;3;4;5;6])
-  [[1;2;3];[4]] (subs 3 ~drop_short:false [1;2;3;4])
+  [[1;2;3];[4]] (subs 3 ~last:CCOpt.return [1;2;3;4])
 *)
 
 let take_while p l =

src/core/CCList.mli

@@ -96,7 +96,7 @@ val partition_map : ('a -> [<`Left of 'b | `Right of 'c | `Drop]) ->
     @since 0.11 *)
 
 val sublists_of_len :
-  ?drop_short:bool ->
+  ?last:('a list -> 'a list option) ->
   ?offset:int ->
   int ->
   'a list ->
@@ -104,11 +104,13 @@ val sublists_of_len :
 (** [sublists_of_len n l] returns sub-lists of [l] that have length [n].
     By default, these sub-lists are non overlapping:
     [sublists_of_len 2 [1;2;3;4;5;6]] returns [[1;2]; [3;4]; [5;6]]
-    @param drop_short if true, last elements are dropped if they do not
-      make a long enough list
     @param offset the number of elements dropped between two consecutive
       sub-lists. By default it is [n]. If [offset < n], the sub-lists
       will overlap; if [offset > n], some elements will not appear at all.
+    @param last if provided and the last group of elements [g] is such
+      that [length g < n], [last g] is called. If [last g = Some g'],
+      [g'] is appended; otherwise [g] is dropped.
+      If [last = CCOpt.return], it will simply keep the last group.
     @raise Invalid_argument if [offset <= 0] or [n <= 0]
     @since NEXT_RELEASE *)