fp.control¶

fp.control expresses loops as functions that pass an accumulator, or carry, from one step to the next. while_loop and fori_loop return the final carry; scan also collects an output from each step. For folds over native iterators, see the iteration tutorial.

Functional loops and scans over native Mojo values: while_loop, fori_loop and scan.

They follow the sequential value semantics of JAX's control-flow operations, without tensor stacking, tracing, autodiff or a JIT. Branch with Mojo's if.

The carry may be any value: a scalar, a SIMD value, a string, a container or a move-only owner. Its type stays fixed, and nothing copies it. A callback is a plain function or a closure, with separate pure and raising overloads, or a shared-receiver library value: Unary for a while_loop body, Binary taking (Int, carry) for a fori_loop body, and a Binary step for scan, whose output type is then named, scan[Y](step, ...). A while_loop predicate may be a BorrowCallable. A call needing keywords or defaults is written as a closure. Callbacks keep one receiver throughout.

Errors stop the loop at once. A predicate and a body must raise one error type, including any embedded origins, or one of them nothing. Nothing retries or rolls back effects.

The compile-time parameter unroll expands that many steps per rolled iteration, with exact bounds checks for the final partial chunk; 1, the default, keeps a rolled loop. unroll=0 fully expands a statically bounded loop: fori_loop[lower, upper, unroll=0](body, initial) or scan[static_length=N, unroll=0](...). Full unrolling without static bounds, a negative factor and an invalid static length are compile errors. while_loop has no unroll parameter.

Summary¶

Functions¶

fori_loop¶

Source: fp/control/loops.mojo

def fori_loop[A, F: def(var Int, var A) -> A, //, unroll: Int = 1](
    lower: Int,
    upper: Int,
    body_fun: F,
    var init_val: A,
) raises Never -> A

Call body_fun(index, carry) for each index in [lower, upper), in ascending order.

Equal or reversed bounds return init_val without a call; negative bounds are valid. Static bounds are written fori_loop[lower, upper](body, carry).

Parameters

• A (Movable & Deinitable, inferred): The carry type: any value, including a move-only owner.
• F (def(var Int, var A) -> A, inferred): The body's type: a plain function or a closure taking the index and the carry.
• unroll (Int): How many steps each rolled iteration expands: 1, the default, keeps a rolled loop, and 0 fully expands a loop with static bounds. Call order and error behavior do not change.

Arguments

• lower (Int): The first index.
• upper (Int): The bound, excluded.
• body_fun (F): The body, borrowed.
• init_val (var A): The first carry, consumed.

Returns

A: The carry after the last call.

Raises

Never: Nothing: Never for a pure body. The raising overloads raise the body's error, which stops the loop at once.

def fori_loop[A, E, F: def(var Int, var A) raises E -> A, //, unroll: Int = 1](
    lower: Int,
    upper: Int,
    body_fun: F,
    var init_val: A,
) raises E -> A

def fori_loop[A, F: Binary, //, unroll: Int = 1](
    lower: Int,
    upper: Int,
    body_fun: F,
    var init_val: A,
) raises F.Error -> A where F.First == Int and F.Second == A and F.Out == A

def fori_loop[
    A, F: def(var Int, var A) -> A, //, lower: Int, upper: Int, unroll: Int = 1,
](
    body_fun: F,
    var init_val: A,
) raises Never -> A

def fori_loop[
    A, E, F: def(var Int, var A) raises E -> A, //, lower: Int, upper: Int,
    unroll: Int = 1,
](
    body_fun: F,
    var init_val: A,
) raises E -> A

def fori_loop[A, F: Binary, //, lower: Int, upper: Int, unroll: Int = 1](
    body_fun: F,
    var init_val: A,
) raises F.Error -> A where F.First == Int and F.Second == A and F.Out == A

scan¶

Source: fp/control/scan.mojo

def scan[
    A, X, Y, F: def(var A, var X) -> Tuple[A, Y], //, unroll: Int = 1,
    static_length: Int = -1,
](
    f: F,
    var init: A,
    var xs: List[X],
    *,
    length: Optional[Int] = None,
    reverse: Bool = False,
) raises ScanError[Never] -> Tuple[A, List[Y]]

Thread a carry through a sequence, collecting one output per element.

The step receives (carry, element) and returns (next_carry, output); the carry and output may have different types, and neither is copied. The initial carry is not emitted (fp.iteration.scan_left emits snapshots instead). Empty input returns the initial carry and an empty list. An iterator input is collected into a List first, so it must be finite, and every pull finishes before the first step. With no input, give length or static_length; the step then receives Tuple[]. In reverse, running sums over [1, 2, 3] give the carry 6 and the outputs [6, 5, 3]. On failure the partial outputs, unused inputs and carry are destroyed.

Parameters

• A (Movable & Deinitable, inferred): The carry type.
• X (Movable & Deinitable, inferred): The element type.
• Y (Movable & Deinitable, inferred): The output type.
• F (def(var A, var X) -> Tuple[A, Y], inferred): The step's type: a plain function or a closure.
• unroll (Int): How many steps each rolled iteration expands: 1, the default, keeps a rolled loop, and 0 fully expands a scan of static_length.
• static_length (Int): The input length, known when the program compiles; -1 when it is not.

Arguments

• f (F): The step, borrowed.
• init (var A): The first carry, consumed.
• xs (var List[X]): The elements, consumed.
• length (Optional[Int]): The input's length, checked against it; it does not truncate.
• reverse (Bool): Whether the steps run from the last element to the first. The outputs keep the positions of their inputs.

Returns

Tuple[A, List[Y]]: The final carry and the outputs, in input order.

Raises

ScanError[Never]: ScanLengthError for a negative, missing or different length, before any step runs.

def scan[
    A, X, Y, E, F: def(var A, var X) raises E -> Tuple[A, Y], //, unroll: Int =
    1, static_length: Int = -1,
](
    f: F,
    var init: A,
    var xs: List[X],
    *,
    length: Optional[Int] = None,
    reverse: Bool = False,
) raises ScanError[E] -> Tuple[A, List[Y]]

def scan[A, X, F: Binary, //, Y, unroll: Int = 1, static_length: Int = -1](
    f: F,
    var init: A,
    var xs: List[X],
    *,
    length: Optional[Int] = None,
    reverse: Bool = False,
) raises ScanError[F.Error] -> Tuple[A, List[Y]] where F.First == A and F.Second == X and F.Out == Tuple[A, Y]

def scan[
    A, X, Y, I: Iterator, F: def(var A, var X) -> Tuple[A, Y], //, unroll: Int =
    1, static_length: Int = -1,
](
    f: F,
    var init: A,
    var xs: I,
    *,
    length: Optional[Int] = None,
    reverse: Bool = False,
) raises ScanError[Never] -> Tuple[A, List[Y]] where I.Element == X

def scan[
    A, X, Y, I: Iterator, E, F: def(var A, var X) raises E -> Tuple[A, Y], //,
    unroll: Int = 1, static_length: Int = -1,
](
    f: F,
    var init: A,
    var xs: I,
    *,
    length: Optional[Int] = None,
    reverse: Bool = False,
) raises ScanError[E] -> Tuple[A, List[Y]] where I.Element == X

def scan[
    A, I: Iterator, F: Binary, //, Y, unroll: Int = 1, static_length: Int = -1,
](
    f: F,
    var init: A,
    var xs: I,
    *,
    length: Optional[Int] = None,
    reverse: Bool = False,
) raises ScanError[F.Error] -> Tuple[A, List[Y]] where F.First == A and F.Second == I.Element and F.Out == Tuple[A, Y]

def scan[
    A, Y, F: def(var A, var Tuple[]) -> Tuple[A, Y], //, unroll: Int = 1,
    static_length: Int = -1,
](
    f: F,
    var init: A,
    *,
    length: Optional[Int] = None,
    reverse: Bool = False,
) raises ScanError[Never] -> Tuple[A, List[Y]]

def scan[
    A, Y, E, F: def(var A, var Tuple[]) raises E -> Tuple[A, Y], //, unroll: Int
    = 1, static_length: Int = -1,
](
    f: F,
    var init: A,
    *,
    length: Optional[Int] = None,
    reverse: Bool = False,
) raises ScanError[E] -> Tuple[A, List[Y]]

def scan[A, F: Binary, //, Y, unroll: Int = 1, static_length: Int = -1](
    f: F,
    var init: A,
    *,
    length: Optional[Int] = None,
    reverse: Bool = False,
) raises ScanError[F.Error] -> Tuple[A, List[Y]] where F.First == A and F.Second == Tuple[] and F.Out == Tuple[A, Y]

while_loop¶

Source: fp/control/loops.mojo

def while_loop[A, //, P: def(A) -> Bool, F: def(var A) -> A](
    cond_fun: P,
    body_fun: F,
    var init_val: A,
) raises _CommonError[Never, Never] -> A where _ErrorCompatible[_CommonError[Never, Never], Never]

Repeat body_fun while cond_fun holds for the carry, and return the carry.

The predicate borrows the carry before each iteration. When it returns False the current carry is returned; otherwise the carry moves into the body, which returns the next one of exactly the same type. Nothing copies the carry.

Parameters

• A (Movable & Deinitable, inferred): The carry type: any value, including a move-only owner.
• P (def(A) -> Bool): The predicate's type: a plain function or a closure.
• F (def(var A) -> A): The body's type: a plain function or a closure.

Arguments

• cond_fun (P): The predicate, borrowed.
• body_fun (F): The body, borrowed.
• init_val (var A): The first carry, consumed.

Returns

A: The carry for which the predicate first returned False.

Raises

Nothing: Never for a pure predicate and body. The raising overloads raise the predicate's or the body's error, which stops the loop at once.

def while_loop[A, P: BorrowCallable, F: Unary](
    cond_fun: P,
    body_fun: F,
    var init_val: A,
) raises _CommonError[P.Failure, F.Error] -> A where P.Payload == A and P.Result == Bool and F.Arg == A and F.Out == A and _ErrorCompatible[_CommonError[P.Failure, F.Error], P.Failure] and _ErrorCompatible[_CommonError[P.Failure, F.Error], F.Error]

def while_loop[A, BE, //, P: def(A) -> Bool, F: def(var A) raises BE -> A](
    cond_fun: P,
    body_fun: F,
    var init_val: A,
) raises _CommonError[Never, BE] -> A where _ErrorCompatible[_CommonError[Never, BE], BE]

def while_loop[A, PE, //, P: def(A) raises PE -> Bool, F: def(var A) -> A](
    cond_fun: P,
    body_fun: F,
    var init_val: A,
) raises _CommonError[PE, Never] -> A where _ErrorCompatible[_CommonError[PE, Never], PE] and _ErrorCompatible[_CommonError[PE, Never], Never]

def while_loop[
    A, PE, BE, //, P: def(A) raises PE -> Bool, F: def(var A) raises BE -> A,
](
    cond_fun: P,
    body_fun: F,
    var init_val: A,
) raises _CommonError[PE, BE] -> A where _ErrorCompatible[_CommonError[PE, BE], PE] and _ErrorCompatible[_CommonError[PE, BE], BE]

Structs¶

ScanLengthError¶

Source: fp/control/scan.mojo

struct ScanLengthError(Copyable, Movable, Writable)

A scan length that is negative, missing or different from the input's; no callback has run.

Implements

Copyable, Writable

Fields

• expected (Int): The length the scan required.
• actual (Int): The length it was given; -1 when no length was given.

ScanStepError¶

Source: fp/control/scan.mojo

struct ScanStepError[E](Movable, Copyable where conforms_to(E, Copyable))

A scan step's error, kept distinct from a length error.

A step that raises StopIteration is reported here too: it is a step error, not exhaustion.

Parameters

• E (Movable & Deinitable): The step's error type.

Implements

Copyable

Fields

• error (E): The step's error, unchanged.

Aliases¶

ScanError¶

Source: fp/control/scan.mojo

comptime ScanError[E] = ScanLengthError if E == Never else Variant[ScanLengthError, ScanStepError[E]]

The error of scan: ScanLengthError alone for a pure step, otherwise a native Variant of it and ScanStepError[E].

Parameters

• E (Movable & Deinitable): The step's error type; Never for a pure step.

Write loops and a scan as expressions¶

"""Functional loops over native Mojo values and sequences."""
from fp.control import while_loop, fori_loop, scan


def below_ten(value: Int) -> Bool:
    return value < 10


def double(var value: Int) -> Int:
    return 2 * value


def add_index(var index: Int, var carry: Int) -> Int:
    return carry + index


def running_total(var carry: Int, var value: Int) -> Tuple[Int, Int]:
    var total = carry + value
    return (total, total)


def main() raises:
    print("while:", while_loop(below_ten, double, 1))
    print("fori:", fori_loop(0, 5, add_index, 0))
    var xs: List[Int] = [1, 2, 3]
    var result = scan(running_total, 0, xs^, reverse=True)
    print("scan:", result[0], result[1])

while: 16
fori: 10
scan: 6 [6, 5, 3]

while_loop(below_ten, double, 1) doubles 1 until the carry is no longer below ten, and returns 16. fori_loop(0, 5, add_index, 0) adds the indices 0 to 4 and returns 10. scan(running_total, 0, xs^, reverse=True) visits 3, 2 and 1, so the carry ends at 6, and returns the outputs in input order, [6, 5, 3].