Mojo Closures

A closure is a function defined inside another function that can capture and use variables from the surrounding scope. Closures let you create small, specialized functions on the fly without writing a full named function at the top level. Mojo supports closures through nested function definitions and the @parameter mechanism for compile-time closures.

The Capture Concept

  Outer scope has variable:  tax_rate = 0.18

  ┌──────────────────────────────────────┐
  │  fn apply_tax(price: Float64):       │
  │                                      │
  │    fn compute() -> Float64:          │  ← inner function (closure)
  │        return price * (1 + tax_rate) │  ← captures price AND tax_rate
  │                                      │
  │    return compute()                  │
  └──────────────────────────────────────┘

  compute() "closes over" both price and tax_rate.
  It carries references to those variables wherever it goes.

Basic Closure

fn make_multiplier(factor: Int) -> fn(Int) -> Int:
    fn multiply(x: Int) -> Int:
        return x * factor   # captures 'factor' from the outer scope
    return multiply

fn main():
    var double = make_multiplier(2)
    var triple = make_multiplier(3)

    print(double(5))    # 10
    print(triple(5))    # 15
    print(double(10))   # 20

Closure factory diagram:

  make_multiplier(2)  →  double  [captures factor=2]
  make_multiplier(3)  →  triple  [captures factor=3]

  double(5):  5 × 2 = 10   (uses its own captured factor)
  triple(5):  5 × 3 = 15   (uses its own captured factor)

Closures as Callbacks

Pass a closure to a function that expects a function argument. This lets you customize the behavior of a general function at the call site.

fn apply_to_each(data: List[Int], func: fn(Int) -> Int) -> List[Int]:
    var result = List[Int]()
    for i in range(len(data)):
        result.append(func(data[i]))
    return result

fn main():
    var numbers = List[Int](1, 2, 3, 4, 5)

    # Closure 1: square each element
    fn square(x: Int) -> Int:
        return x * x

    # Closure 2: add 10 to each element
    fn add_ten(x: Int) -> Int:
        return x + 10

    var squared = apply_to_each(numbers, square)
    var shifted = apply_to_each(numbers, add_ten)

    for i in range(len(squared)):
        print(squared[i], end=" ")   # 1 4 9 16 25
    print("")
    for i in range(len(shifted)):
        print(shifted[i], end=" ")   # 11 12 13 14 15
    print("")

Capturing Mutable State

A closure can capture a variable and modify it across multiple calls, building a form of stateful behavior.

fn make_counter() -> fn() -> Int:
    var count = 0

    fn increment() -> Int:
        count += 1
        return count

    return increment

fn main():
    var counter = make_counter()
    print(counter())   # 1
    print(counter())   # 2
    print(counter())   # 3

    # A second counter has its own independent count
    var other = make_counter()
    print(other())    # 1  ← starts fresh
    print(counter())  # 4  ← original continues

State isolation diagram:

  counter [count=0] → call → count=1 → call → count=2 → ...
  other   [count=0] → call → count=1   (separate captured variable)

@parameter Closures (Compile-Time)

When you mark a nested function with @parameter, Mojo evaluates it at compile time. This is the mechanism used by vectorize, parallelize, and tile throughout the standard library.

from algorithm import vectorize
from memory import UnsafePointer

fn scale_all(data: UnsafePointer[Float32], n: Int, factor: Float32):
    @parameter
    fn scale_chunk[width: Int](i: Int):
        var v = SIMD[DType.float32, width].load(data + i)
        (v * factor).store(data + i)

    vectorize[scale_chunk, 8](n)

The inner scale_chunk function captures data and factor from the outer scope. The @parameter decorator signals that the compiler should inline and specialize this closure for each SIMD width.

Closures for Sorting

Pass a closure as a comparator to implement custom sort orders without writing a separate named function.

fn sort_by[compare: fn(Int, Int) -> Bool](inout arr: List[Int]):
    var n = len(arr)
    for i in range(n):
        for j in range(0, n - i - 1):
            if compare(arr[j + 1], arr[j]):   # swap if compare says so
                var temp = arr[j]
                arr[j] = arr[j + 1]
                arr[j + 1] = temp

fn main():
    var nums = List[Int](5, 2, 8, 1, 9, 3)

    fn ascending(a: Int, b: Int) -> Bool:
        return a < b

    fn descending(a: Int, b: Int) -> Bool:
        return a > b

    sort_by[ascending](nums)
    for i in range(len(nums)):
        print(nums[i], end=" ")   # 1 2 3 5 8 9
    print("")

    sort_by[descending](nums)
    for i in range(len(nums)):
        print(nums[i], end=" ")   # 9 8 5 3 2 1
    print("")

Closures vs Named Functions

Named function:               Closure:
  fn square(x: Int) -> Int:    fn main():
      return x * x                 fn square(x: Int) -> Int:
                                       return x * x
  Visible everywhere.              Visible only inside main().
  Cannot capture outer vars.       Can capture outer vars.
  Reusable across files.           Single-use, contextual.

Key Takeaways

A closure is a nested function that captures variables from the enclosing scope. Each closure instance carries its own copy of the captured variables, enabling independent state. Pass closures as arguments to create flexible, reusable higher-order functions. Use @parameter closures for compile-time specialization inside vectorize, parallelize, and tile. Closures make sort comparators, event callbacks, and factory patterns concise and readable. Mojo function types are written as fn(ArgType) -> ReturnType and can be stored in variables or passed as parameters.