refactor: use map function to implement map block (#1096)

Signed-off-by: Louis Mandel <[email protected]>

Author: mandel

src/pdl/pdl_interpreter.py

@@ -12,6 +12,7 @@
 import warnings
 from asyncio import AbstractEventLoop
 from functools import partial, reduce
+from itertools import count
 from os import getenv
 
 warnings.filterwarnings("ignore", "Valid config keys have changed in V2")
@@ -21,6 +22,7 @@
     Any,
     Generator,
     Generic,
+    Iterable,
     Optional,
     Sequence,
     Tuple,
@@ -952,50 +954,48 @@ def process_block_body(
                 yield_result(result.result(), block.kind)
             trace = block.model_copy(update={"pdl__trace": iter_trace})
         case MapBlock():
-            results = []
             background = DependentContext([])
-            iter_trace = []
             iteration_state = state.with_yield_result(False)
             block, items, length = _evaluate_for_field(scope, block, loc)
             block, max_iterations = _evaluate_max_iterations_field(scope, block, loc)
             block = _evaluate_join_field(scope, block, loc)
             map_loc = append(loc, "map")
             iidx = 0
             try:
-                saved_background = IndependentContext([])
-                while True:
+                if max_iterations is not None:
+                    index_iterator: Any = range(max_iterations)
+                else:
+                    index_iterator = count()
+                if items is not None and length is not None:
+                    items_iterator = (
+                        {k: elems[i] for k, elems in items.items()}
+                        for i in range(length)
+                    )
+                else:
+                    items_iterator = ({} for _ in count())
+
+                def loop_body(iidx, items):
                     iteration_scope = scope_init
                     if block.index is not None:
                         iteration_scope = iteration_scope | {block.index: iidx}
-                    if max_iterations is not None and iidx >= max_iterations:
-                        break
-                    if length is not None and iidx >= length:
-                        break
-                    iteration_state = iteration_state.with_iter(iidx)
-                    if items is not None:
-                        for k in items.keys():
-                            iteration_scope = iteration_scope | {k: items[k][iidx]}
-                    (
-                        iteration_result,
-                        iteration_background,
-                        iteration_scope,
-                        body_trace,
-                    ) = process_block(
-                        iteration_state,
+                    iteration_scope = iteration_scope | items
+                    return process_block(
+                        iteration_state.with_iter(iidx),
                         iteration_scope,
                         block.map,
                         map_loc,
                     )
-                    saved_background = IndependentContext(
-                        [saved_background, iteration_background]
-                    )
-                    results.append(iteration_result)
-                    iter_trace.append(body_trace)
-                    iteration_state = iteration_state.with_pop()
-                    iidx = iidx + 1
+
+                # with ThreadPoolExecutor(max_workers=4) as executor:
+                #     map_output = executor.map(
+                map_output = map(  # pylint: disable=bad-builtin
+                    loop_body, index_iterator, items_iterator
+                )
+                results, _, _, traces = _split_map_output(map_output)
+                # saved_background = IndependentContext(backgrounds)
             except PDLRuntimeError as exc:
-                iter_trace.append(exc.pdl__trace)
-                trace = block.model_copy(update={"pdl__trace": iter_trace})
+                traces = [exc.pdl__trace]  # type: ignore
+                trace = block.model_copy(update={"pdl__trace": traces})
                 raise PDLRuntimeError(
                     exc.message,
                     loc=exc.loc or map_loc,
@@ -1005,7 +1005,7 @@ def process_block_body(
             # background = saved_background  # commented because the block do not contribute to the background
             if state.yield_result and not iteration_state.yield_result:
                 yield_result(result.result(), block.kind)
-            trace = block.model_copy(update={"pdl__trace": iter_trace})
+            trace = block.model_copy(update={"pdl__trace": traces})
         case ReadBlock():
             result, background, scope, trace = process_input(state, scope, block, loc)
             if state.yield_result:
@@ -1065,6 +1065,21 @@ def process_block_body(
     return result, background, scope, trace
 
 
+def _split_map_output(
+    map_output: Iterable[Tuple[PdlLazy[Any], LazyMessages, ScopeType, BlockType]],
+) -> Tuple[list[PdlLazy[Any]], list[LazyMessages], list[ScopeType], list[BlockType]]:
+    results = []
+    backgrounds = []
+    scopes = []
+    traces = []
+    for result, background, scope, trace in map_output:
+        results.append(result)
+        backgrounds.append(background)
+        scopes.append(scope)
+        traces.append(trace)
+    return results, backgrounds, scopes, traces
+
+
 BlockTVarEvalFor = TypeVar("BlockTVarEvalFor", bound=RepeatBlock | MapBlock)