Mojo By Example: A Comprehensive Introduction to the Mojo Programming Language

14. Parameters and compile-time programming

14.1. Compile-time execution of code

The programs we write are compiled by the compiler and then an executable file is created out of it. This executable file is later run on an operating system. So the code we write is executed when we run the program. So if we write a loop, then that loop gets executed when the program is run. Let’s imagine that we need to execute a function for which the inputs are well known already at the time of writing the code. Such a function can be executed at compile-time, and the results kept in the executable file. This means that when the program is run, it just takes the pre-calculated value, saving valuable CPU time during the execution.

Unlike many mainstream languages, Mojo allows execution of code at compile-time. Mojo even has very few restrictions on what kind of code can be executed at compile-time.

We already saw the use of alias to store a constant. When we assign a function’s value to an alias constant, Mojo executes that function at compile-time and assigns the result as a constant value to the alias.

For example see the following code listing.

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

alias added = add(1, 2)

Usage:

    print(added)

14.2. Conditional execution at compile-time

We can conditionally call a function, or define a constant. That is, the if statement also can be used at compile time. Mojo allows us to pass an environment variable at compile time, which we can access and conditionally compile our code. To pass an environment variable to the compiler, you can use the option -D <var>=<val>. For example, mojo -D add_it=True filename.mojo.

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

# Execute with "mojo -D add_it=True <filename>" for a non-zero value
from sys import is_defined
alias added_conditionally = add(1, 2) if is_defined["add_it"]() else 0

Usage:

    print(added_conditionally)

14.3. Parameters in functions

We saw that we can use alias to execute functions at compile time. What if we do not want to execute the function at compile-time, but just "parameterize" the function at compile-time, so that when it is finally run during the program execution, it uses that passed-in parameters? Mojo provides a solution to that by providing capability for compile-time parameterization. To achieve this, we pass parameters within square brackets. For example: add[my_param: Int](). So in short, when we want to pass to a function values at runtime, we pass it within parenthesis, and if we want to pass to a function values at compile-time, we pass it within square brackets. Note that we are distinguishing between parameters and arguments. In Mojo parameters are compile-time values and arguments are runtime values.

fn add[cond: Bool](a: Int, b: Int) -> Int: 
    return a + b if cond else 0

Usage:

    # Execute with "mojo -D add_it=True <filename>" for a non-zero value
    from sys import is_defined
    print(add[is_defined["add_it"]()](3, 4))

In the previous example, we are passing a boolean value as a parameter to the add function. The value of the boolean is taken from a compile-time command line argument add_it using the Mojo option -D. The compiler will hardcode True in the parameter cond of the function add if the add_it was set (the value of the add_it is irrelevant here). Irrespective of where and how the compiled file is executed, the cond will have a constant value set at the compile-time.

In addition to values being passed as parameters to functions, we can also pass types themselves. If a concrete struct is expected in the parameter, then Mojo expects a value to be passed at the function call site. However, if a trait is expected in the parameter, then Mojo expects a type to be passed at the function call site.

In the previous example we saw that add expected a value of type Bool. In the next example, the add_ints expects a trait Intable as a parameter. That trait is then later used to declare the types of the function arguments. At the function call site, you can see that Int is being passed as parameter.

fn add_ints[IntType: Intable](a: IntType, b: IntType) -> Int: 
    return int(a) + int(b)

Usage:

    print(add_ints[Int](3, 4))

14.4. Keyword parameters

So far we have seen how we can pass parameters by position. Similar to the keyword arguments that we pass to functions, we can pass parameter values using the name of the parameter. The rules of the keyword parameters are the same as the keyword arguments for functions.

fn div_compile_time[a: Int, b: Int]() -> Float64: 
    return a / b

Usage:

    print(div_compile_time[b=3, a=4]())

The previous example shows that we can pass parameter values with the param names, in which case the order of the parameters is not relevant.

14.5. Inferred-only parameters

Mojo allows parameter types to depend on other parameter types. For example, suppose we have two structs, Scheme and Location. We could define Location as a struct that takes a Scheme value as its input parameter, as in Location[scheme: Scheme]. This means that Location depends on Scheme.

struct Scheme:

    alias HTTP = Scheme("http")
    alias FTP = Scheme("ftp")

    var scheme: String

    fn __init__(inout self, scheme: String):
        self.scheme = scheme

    fn __str__(self) -> String:
        return self.scheme

struct Location[scheme: Scheme]:
    
    var location: String

    fn __init__(inout self, location: String):
        self.location = location
    
    fn __str__(self) -> String:
        return str(scheme) + "://" + self.location

Suppose that we now define a function that uses Location. We now need to also declare Scheme parameter as otherwise the compiler does not know what the Location input parameter scheme means.

fn print_location[scheme: Scheme, location: Location[scheme]]():
    print(str(location))

This has an unfortunate impact on the ergonomics of the usage of the function, as now the caller has to specify both the Scheme and Location with again the same Scheme value. This is an unnecessary duplication.

    print_location[Scheme.FTP, Location[Scheme.FTP]("r.net")]()

Mojo provides a solution for this. Similar to declaration of positional-only and keyword-only function arguments, Mojo provides a syntax for "inferred-only" parameters using // as the delimiter. All the parameters that are expected to be inferred will appear before the // delimiter. Those parameters are not to be passed by the caller, instead they would be automatically inferred by the compiler based on their usage in the following parameters.

fn print_location2[scheme: Scheme, //, location: Location[scheme]]():
    print(str(location))

Usage:

    print_location2[Location[Scheme.FTP]("r.net")]()

Here we have to provide Scheme.FTP only once as the parameter scheme: Scheme will get automatically inferred.

It is also possible to have just inferred-only parameters, while the inference is happening within the function arguments.

fn print_location3[scheme: Scheme, //](location: Location[scheme]):
    print(str(location))

Usage:

    print_location3(Location[Scheme.FTP]("r.net"))

14.6. Variadic parameters

Sometimes we want to be able to pass any number of parameters, without being restricted to a particular number of parameters. When we prefix a parameter with *, Mojo allows us to pass any number of values to it.

The following is an example of positional variadic parameters:

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

Usage:

    print(add_all[1, 2, 3]())

In the above example, we can see that at the function call site, we can pass any number of parameters since the function definition prefixes its parameter a with a *. The function definition then iterates over the a after wrapping it in a VariadicList and calculates the sum.

14.7. Default values in parameters

Mojo allows default values to be used for parameters.

fn sub[cond: Bool = False](a: Int, b: Int) -> Int: 
    return a + b if cond else 0

Usage:

    print(sub(3, 4)) # Default value is taken
    print(sub[True](3, 4)) # Override the default value

In the example, we have assigned a default value for parameter cond, which allow us to make the call to sub without passing any value to the parameter cond, effectively making the parameter cond an optional parameter. This makes for ergonomic APIs using sensible defaults wherever it is possible.

14.8. Parameters in structs, traits

Similar to functions, we can also pass compile-time parameters to structs, traits.

struct MyStruct[T: Intable, cond: Bool]:
    var value: Int

    fn __init__(inout self, value: T):
        self.value = int(value)

    fn get_value(self) -> Int:
        return self.value if cond else 0

Usage:

    print(MyStruct[Int, True](10).get_value())
    print(MyStruct[Float16, False](11.5).get_value())

We can also have keyword parameters for structs, and traits.

    print(MyStruct[cond=True, T=Float32](2.5).get_value())

In the previous example, the parameters were passed and processed similar to how we did in functions. Basically, what we can do with parameters for functions, we can do the same for structs, and traits.

14.9. Conditional conformance

We can also overload methods based on concrete types of self itself. In the following code listing, the version of the hello method chosen depends on the concrete type of the passed parameter T. Note that here the self type is explicitly defined, unlike how we usually specify self.

struct HelloMessage[T: AnyType]:
    
    fn __init__(inout self): ...

    fn hello(self: HelloMessage[Int]): 
        print("Hello 10")

    fn hello(self: HelloMessage[Float16]): 
        print("Hello 10.5")

Usage:

    HelloMessage[Int]().hello() # Prints 'Hello 10'
    HelloMessage[Float16]().hello() # Prints 'Hello 10.5'

You can also use numeric parameters.

struct HelloCount[count: Int]:
    
    fn __init__(inout self): ...

    fn hello(self: HelloCount[2]): 
        print("Count was 2")

    fn hello(self: HelloCount[3]): 
        print("Count was 3")

Usage:

    HelloCount[2]().hello() # Prints 'Count was 2'
    HelloCount[3]().hello() # Prints 'Count was 3'
    # HelloCount[4]().hello() # Uncommenting this would result in error as there is no such overload

Conditional conformance also works with init method.

14.10. Custom compile-time checks

When we develop a program, we often make assumptions about the arguments we receive or the context in which we execute a function and so on. We can use if statements to validate those assumptions, but it is possible that there is a performance cost to such validations. Many programming languages provide a facility known as assertion, to validate those assumptions with minimal impact to the runtime performance of the code.

Mojo goes one step further by providing compile-time assertions with the function constrained.

fn print_times[times: Int]():
    constrained[times > 0, "times must be greater than zero"]()
    for i in range(times):
        print(i)

fn main():
    print_times[2]()
    print_times[0]()

If you try to compile the code listed above, you would get a compile time error, with the message: times must be greater than zero.

This happens because in our function print_times we have an assertion using constrained function that checks that our compile-time parameter times is greater than zero. If we call the function with a value for times that is greater than zero, then the code compiles without any errors. However, if we pass a value that is less than or equal to zero, it will produce the same compile time error message as the one that is passed as the second parameter of the constrained function call.

The constrained function is quite useful to validate our assumptions about given parameters at compile-time. If during the compile-time the constrained function executes successfully, then that piece of validation code does not even have to appear in the final binary, resulting in zero performance impact at runtime.