14. Advanced usage of functions

In previous chapters you saw how we declare types of variables and how we use them. In Mojo you can assign also a function to a variable. The type of the variable is determined by the function’s signature, i.e., it is a combination of the argument types and return type of the function.

In the following code listing, you can see the variable my_fn_var is of type fn(Int, Int) → Int and is assigned a function with the same signature as the type of my_fn_var. Another thing to note is that the assigned function does not have the trailing () which is usual in a function call. This is because we are not calling the function adder, instead we are binding the function to the variable my_fn_var. In fact, we do not want to call adder at that point of time.

fn adder(a: Int, b: Int) -> Int: 
    return a + b

Usage:

    var my_fn_var: fn(Int, Int) -> Int = adder
    print(my_fn_var(4, 3))

In the example, we see that we defined a function adder and then assigned that function to a variable. We later execute that variable as if it is a function.

14.1. Higher-order functions

The ability to assign a function to a variable allows us to implement some interesting use cases. We can pass a function as an argument to another function. We can also return a function as the result from another function. A function that can take a function as an argument, or can return a function as the result is called a higher-order function.

There are some interesting uses of higher-order functions, such as the ability to define generic functions that takes any function as an argument and executes the given function based on some condition, or inside a loop, and so on.

fn adder(a: Int, b: Int) -> Int: 
    return a + b

fn suber(a: Int, b: Int) -> Int: 
    return a - b

fn exec(x: Int, y: Int, bin_op: fn(Int, Int) -> Int) -> Int:
    var result: Int = 0
    for i in range(10):
        result += bin_op(x, y)
    return result

Usage:

    print(exec(10, 5, adder))
    print(exec(10, 5, suber))

In this code listing we see that the function exec is a higher-order function that takes another function as argument. The exec function executes the passed-in function within a loop to calculate a result. We can pass two different functions to the same exec function, and the exec will treat both the passed functions the same way. In this way we have built a generic function that does not need to know what the passed-in does, instead it just executes them and calculates results.

Functions can also passed as parameters to another function. The main difference is that the functions are then passed at compile-time instead of at runtime.

fn diver(a: Int, b: Int) -> Float16: 
    return a / b

fn exec_param[bin_op: fn(Int, Int) -> Int](x: Int, y: Int) -> Int:
    var result: Int = 0
    for i in range(10):
        result += bin_op(x, y)
    return result

Usage:

    print(exec_param[adder](10, 5))
    print(exec_param[suber](10, 5))

14.2. Closure

So far we have seen that we can pass arguments to functions and they would use those arguments within their body. There is another technique for functions to get values from outside of the function body, which is commonly known as closure.

In a closure, we define a function that captures values outside of its function body. The values that are captured must be defined before the definition of the function itself. Another constraint is that the data type of the captured values must implement __copyinit__, as the value of the variable is copied over to the function.

fn exec_rt_closure(x: Int, bin_op_cl: fn(Int) escaping -> Int) -> Int:
    var result: Int = 0
    for i in range(10):
        result += bin_op_cl(x)
    return result

Usage:

    var rt_y: Int = 5
    fn ander(x: Int) -> Int: 
        return x & rt_y
    print(exec_rt_closure(12, ander))

The closure shown above is known as a runtime closure. The type of the closure is fn() escaping → T. Note that the captured values are owned by the closure. Runtime closure can be passed as argument to other functions.

However, the runtime closure cannot be passed as a parameter to other functions. In order to pass a closure as a parameter to other functions, we need to use a compile-time closure. Such closures are decorated by @parameter. The type of the closure is fn() capturing → T. The following example demonstrates such a compile-time closure.

fn exec_ct_closure[bin_op_cl: fn(Int) capturing -> Int](x: Int) -> Int:
    var result: Int = 0
    for i in range(10):
        result += bin_op_cl(x)
    return result

Usage:

    var ct_y: Int = 10
    @parameter
    fn multer(x: Int) -> Int: 
        return x * ct_y
    print(exec_ct_closure[multer](10))

14.3. Variadic function

You have already seen how the built-in print function is able to take any number of arguments. This is possible because of Mojo’s support for variadic functions. The variadic argument is prefixed by * for positional arguments and ** for keyword arguments. Data type of positional variadic argument is VariadicList[T] and of keyword variadic argument is Dict[K, V].

The following is an example of positional variadic arguments:

fn add_all(*a: Int) -> Int: 
    var result: Int = 0
    for i in a:
        result+= i
    return result

Usage of positional variadic arguments:

    print(add_all(1, 2, 3, 4, 5))

The following is an example of keyword variadic arguments:

fn names_dob(**namedobs: String) raises: 
    for name in namedobs.keys():
        print(name[], namedobs[name[]])

Usage of keyword variadic arguments:

    names_dob(ik="14/4/72", pk="15/5/81", ani="16/3/23", ad="22/6/17")

In the examples shown, the add_all has exactly one argument prefixed by *. The argument a is then iterated over and its elements extracted to calculate the sum. In case of the function names_dob, the argument namedobs is iterated over to print both the key name and the value associated with the key. Note that the key of the Dict is a Reference and therefore needs to be dereferenced.

Variadic arguments makes it possible to provide easy-to-use APIs by allowing any number of arguments to be passed directly to the function instead of having to wrap them up in a list or dictionary.

14.4. Overloading

There is a saying among programmers that naming is hard. For example, a function to add two numbers would be best called add. However, we may need to add two integers, one integer and a float, and so on. In some programming languages, that would have resulted in convoluted names such as add_ints, add_int_float and so on. Thankfully, Mojo provides a feature called function overloading. This feature allows us to define the same name to multiple functions as long as their argument types, parameter types or number of arguments are different. If two functions have the same argument types but have different result types, Mojo would complain as overloading is supported only for arguments and parameter types.

The following is an example of overloading by different argument types:

fn add(a: Int, b:Int) -> Int:
    return a + b

fn add(a: Int, b:Float16) -> Int:
    return int(a + b)

Usage of the overloaded functions:

    print(add(1, 2))
    print(add(3, 2.4))

The following is an example of overloading by different parameter types:

fn add[a: Int, b:Int]() -> Int:
    return a + b

fn add[a: Bool, b: Bool]() -> Int:
    var ai: Int = 1 if a else 0
    var bi: Int = 1 if b else 0
    return ai + bi

Usage of the overloaded functions:

    print(add[Int(1), Int(2)]())
    print(add[True, False]())

The following is an example of overloading by different number of arguments, even with the same type of arguments:

fn sub(a: Int, b:Int) -> Int:
    return a - b

fn sub(a: Int, b:Int, c:Int) -> Int:
    return a - b - c

Usage of the overloaded functions:

    print(sub(1, 2))
    print(sub(1, 2, 3))

← Previous: Parameters and compile-time programming | ↑ Up: Mojo By Example: A Comprehensive Introduction to the Mojo Programming Language | Next: Python interoperability →