Introduce the 'flat_map' built-in method in the ''Iterable' type

Author: HugoGGuerrier

lkql_jit/language/src/main/java/com/adacore/lkql_jit/built_ins/methods/IterableMethods.java

@@ -121,4 +121,26 @@ protected BaseLKQLLazyList onIterable(Iterable iterable) {
             return new LKQLFlattenResult(iterable);
         }
     }
+
+    @BuiltInMethod(
+        name = "flat_map",
+        doc = """
+        Given an iterable and a function that takes one argument and return another iterable \
+        value, return a new iterable, result of the function application on all elements, flatten \
+        in a sole iterable value. The returned iterable value is lazy."""
+    )
+    abstract static class FlatMapExpr extends BuiltInBody {
+
+        @Specialization(
+            limit = Constants.SPECIALIZED_LIB_LIMIT,
+            guards = "function.parameterNames.length == 1"
+        )
+        protected BaseLKQLLazyList onIterable(
+            Iterable iterable,
+            LKQLFunction function,
+            @CachedLibrary("function") InteropLibrary functionLibrary
+        ) {
+            return new LKQLFlattenResult(new LKQLMapResult(iterable, function, functionLibrary));
+        }
+    }
 }

testsuite/tests/interpreter/builtins/flat_map/script.lkql

@@ -0,0 +1,15 @@
+# Test a simple flat map
+val lst = [1, 2, 3]
+val fm_1 = lst.flat_map((x) => [x * x, x.img])
+print(fm_1.to_list)
+
+# Ensure flat map is lazy
+val fm_2 = lst.flat_map((x) => { print("Processing " & x.img); [x * x, x.img] })
+print(fm_2[1].img)
+print(fm_2[2].img)
+print(fm_2[5].img)
+print(fm_2[4].img)
+
+# Check that function not returning iterables triggers an error
+val fm_3 = lst.flat_map((x) => x*2)
+print(fm_3[1])

testsuite/tests/interpreter/builtins/flat_map/test.out

@@ -0,0 +1,9 @@
+[1, "1", 4, "2", 9, "3"]
+Processing 1
+1
+"1"
+Processing 2
+Processing 3
+9
+"2"
+error: Type error: expected Iterable but got Int

testsuite/tests/interpreter/builtins/flat_map/test.yaml

@@ -0,0 +1,3 @@
+driver: 'interpreter'
+project: 'default_project/default.gpr'
+lkt_refactor: True

user_manual/generated/std.rst

@@ -173,6 +173,10 @@ Methods for `LazyList`
 
     Return the content of the iterable object with each element associated to its index in a tuple: [(<index>, <elem>), ...]
 
+.. method:: LazyList.flat_map(this, function)
+
+    Given an iterable and a function that takes one argument and return another iterable value, return a new iterable, result of the function application on all elements, flatten in a sole iterable value. The returned iterable value is lazy.
+
 .. method:: LazyList.flatten(this)
 
     Given an iterable of iterables, flatten all of them in a resulting iterable value. The returned value is lazy.
@@ -220,6 +224,10 @@ Methods for `List`
 
     Return the content of the iterable object with each element associated to its index in a tuple: [(<index>, <elem>), ...]
 
+.. method:: List.flat_map(this, function)
+
+    Given an iterable and a function that takes one argument and return another iterable value, return a new iterable, result of the function application on all elements, flatten in a sole iterable value. The returned iterable value is lazy.
+
 .. method:: List.flatten(this)
 
     Given an iterable of iterables, flatten all of them in a resulting iterable value. The returned value is lazy.
@@ -543,6 +551,10 @@ Methods for `SelectorList`
 
     Return the content of the iterable object with each element associated to its index in a tuple: [(<index>, <elem>), ...]
 
+.. method:: SelectorList.flat_map(this, function)
+
+    Given an iterable and a function that takes one argument and return another iterable value, return a new iterable, result of the function application on all elements, flatten in a sole iterable value. The returned iterable value is lazy.
+
 .. method:: SelectorList.flatten(this)
 
     Given an iterable of iterables, flatten all of them in a resulting iterable value. The returned value is lazy.