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

12. Parameters and compile-time programming

12.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 comptime to store a constant. When we assign a function’s value to an comptime constant, Mojo executes that function at compile-time and assigns the result as a constant value to the comptime.

For example see the following code listing.

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

comptime added = add(1, 2)

Usage:

    print(added)

12.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.

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

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

Usage:

    print(added_conditionally)

12.3. Compile-time `if`

You can use comptime if to execute conditional code at compile-time. The comptime if is similar to the regular if, but it is executed at compile-time. The code inside the comptime if block is only compiled if the condition is true at compile-time. When comptime is used for if statement, it also applies to the else and elif blocks.

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

12.4. Compile-time `for`

Similar to comptime if, we can also use comptime for to execute a loop at compile-time. The code inside the comptime for is unrolled at compile-time, and the results are kept in the executable file. This means that actually the loop does not exist at runtime, and there is no performance cost for it at runtime.

    comptime for i in range(5):
        print(i)

The above code approximately compiles to a code similar to the following code listing, where the loop is unrolled and the results are kept in the executable file.

    print(0)
    print(1)
    print(2)
    print(3)
    print(4)

12.5. Compile-time expressions

If you want to execute an expression at compile-time, you can use comptime for that. The expression will be evaluated at compile-time, and the result will be kept in the executable file. This avoids having to create a comptime constant or a function just to evaluate that expression.

    comptime list = [1, 2, 3]
    print(comptime(len(list))) # Prints 3

12.6. Parameters in functions

We saw that we can use comptime 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.

def 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 std.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.

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

Usage:

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

12.7. 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.

def div_compile_time[a: Int, b: Int]() -> Float64: 
    return Float64(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.

12.8. 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(ImplicitlyCopyable):

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

    var scheme: String

    def __init__(out self, scheme: String):
        self.scheme = scheme

    def as_string(self) -> String:
        return self.scheme

struct Location[scheme: Scheme](ImplicitlyCopyable):
    
    var location: String

    def __init__(out self, location: String):
        self.location = location
    
    def as_string(self) -> String:
        return Self.scheme.as_string() + "://" + 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.

def print_location[scheme: Scheme, location: Location[scheme]]():
    print(location.as_string())

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.

def print_location2[scheme: Scheme, //, location: Location[scheme]]():
    print(location.as_string())

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.

def print_location3[scheme: Scheme, //](location: Location[scheme]):
    print(location.as_string())

Usage:

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

12.9. 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:

def add_all[*a: Int]() -> Int: 
    var result: Int = 0
    for i in 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.

12.10. Default values in parameters

Mojo allows default values to be used for parameters.

def 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.

12.11. Parameters in structs

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

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

    def __init__(out self, value: Self.T):
        self.value = Int(value)

    def get_value(self) -> Int:
        return self.value if Self.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.

    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.

You can also use comptime to define an alias for a struct with specific parameters. This is useful when you want to reuse the same struct with the same parameters in multiple places in your code.

    comptime MyStructAlias = MyStruct[Int, True]
    print(MyStructAlias(20).get_value())

12.12. Trait composition

We saw how a struct can implement multiple traits. We also saw how a trait can inherit from multiple traits. In order to achieve good code reuse, it is adviced to define traits that are focused on a single responsibility. This is called the Interface Segregation Principle in software design. When we strictly follow this principle, we end up with many small traits that define a single behavior. Now suppose that we want to define a reusable method that uses multiple traits. One way to achieve this is to define a new trait that inherits from the individual traits, and then define the reusable method to use that trait. However, this approach results in a proliferation of traits that are just used to compose other traits.

Mojo provides with an alternative approach to achieve the same goal without the need to define a new trait. This is called trait composition. We can define parameters to a function that are a composition of multiple traits. This is done by using the & operator between the traits.

trait Flyable:
    def fly(self) -> String: ...

trait Walkable:
    def walk(self) -> String: ...

struct Crow (Flyable, Walkable):
    
    def __init__(out self): ...

    def fly(self) -> String:
        return "Crow is flying"
    
    def walk(self) -> String:
        return "Crow is walking"

def bird[T: Flyable & Walkable](animal: T):
    print(animal.fly())
    print(animal.walk())

def main():
    var crow = Crow()
    bird(crow)

Trait composition also works with comptimes. For example:

comptime Bird = Flyable & Walkable

Sharp eyed readers would have noticed similarity between trait composition and intersection types in some programming languages.

12.13. 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]:
    
    def __init__(out self): ...

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

    def 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]:
    
    def __init__(out self): ...

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

    def 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.

12.14. `where` clauses

Mojo provides where clauses to provide additional constraints on parameters, which are not possible to express using just the type system. The where clause expressions are evaluated at compile-time. If the expression evaluates to False, then a compile-time error is produced.

struct PositiveNumber[value: Int where value >= 0]:
    def __init__(out self):
        ...
    def print_value(self):
        print("The positive number is: " + String(self.value))

Usage:

    var num = PositiveNumber[5]()
    num.print_value()

    #var num2 = PositiveNumber[-3]()  # Uncommenting this will cause a compile-time error

Here we have defined a struct PositiveNumber that has a where clause to ensure that the value passed to it is greater than zero. If we try to create an instance of PositiveNumber with a value that is lesser than zero, we get a compile-time error.

where clauses help developers and especially library authors to define stronger contracts on their APIs, ensuring that the users of those APIs have clear guidance on what the boundaries of the APIs are.

12.15. 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 comptime assert.

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

def 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 comptime assert 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 argument of the comptime assert statement.

The comptime assert statement is quite useful to validate our assumptions about given parameters at compile-time. If during the compile-time the comptime assert statement 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.

12.16. Parameterized errors

Mojo allows us to define parameterized errors. This is useful when we want to abstract away the type of error that could be raised, and let the caller decide on what error type to use.

trait ErrorTrait(Writable):
    def __init__(out self, message: String): ...

struct MyError(ErrorTrait): 
    var message: String

    def __init__(out self, message: String):
        self.message = message

def raise_error[ErrorType: ErrorTrait](cond: Bool) raises ErrorType:
    if cond:
        raise ErrorType("Provided condition is True")
    else:
        print("No error")

Usage:

    try:
        raise_error[MyError](True)
    except e:
        print("Error raised:", e)

In the above example, we defined a custom Trait that represents an error. It requires an init method that takes a String message. The function raise_error takes a parameter ErrorType that is constrained to be a type that implements the ErrorTrait. Inside the function, we raise an instance of the parameterized type as error. At the function call site, we specify which concrete error type to use, in this case a struct MyError that implements the ErrorTrait.

We have basically allowed the caller to determine what error type to use, as long as it conforms to the ErrorTrait. This is quite useful when we want to write generic libraries that can be used in different contexts with different error types.

12.17. Partially bound parameters

In some cases, we may want to partially bind parameters to a function, and leave the rest of the parameters to be specified by the caller.

    comptime IntDict = Dict[Int, _] # The key type is bound but not the value type.
    # var a = IntDict() # Uncommenting this line will result in a compile time error because the type of the values is not specified.
    var b = IntDict[String]() # Now the type of the values is fully specified as String.
    b[1] = "Hello"
    b[2] = "World"
    print(b) # Prints: {1: Hello, 2: World}

In the above example, we have defined a dictionary IntDict that takes an Int as key and a value of an unknown type (represented by _). We have partially bound the type of the key to be Int, while leaving the type of the value to be specified by the caller. This allows us to create different dictionaries with different value types, while still having the key type fixed as Int. Here you would have noticed that we could reuse Dict type while keeping the type of the key fixed.

You can treat partially bound types similar to other types. For example, you can create an instance of a partially bound type, and pass it to a function that expects that type.

    def use_int_dict(d: IntDict[_]): # You can use partially bound types similar to other types.
        ...
    use_int_dict(b)

Partially binding parameters is quite useful to create reusable abstractions with some fixed parameters and some flexible parameters. Such types come in handy when we want to create libraries that can be used in different contexts with different types, while still having some fixed parameters that are common across those contexts.

12.18. Parametric aliases

Earlier we saw that we can use comptime to define an alias for a struct with specific parameters. We can also define aliases with parameters.

def main():
    # tag::parametric_alias[]
    comptime IntDict[T: Intable & Copyable] = Dict[Int, T] # This defines a parametric alias for a dictionary where the key type is always Int and the value type is a parameter T that is to be provided.
    var a = IntDict[Int]() # We can use the alias similar to a normal type, but the value type needs to satisfy the constraints specified in the alias definition (in this case, it needs to implement the traits Intable and Copyable).
    a[1] = 1
    a[2] = 2
    print(a) # Prints: {1: 1, 2: 2}
    # var b = IntDict[String]() # Uncommenting this line will result in a compile time error because the value type String does not satisfy the constraints specified in the alias definition (it does not implement the traits Intable and Copyable).
    # end::parametric_alias[]

In the above example, we have defined a parametric alias StringKeyedDict for a dictionary where the key type is always String and the value type is a parameter T that is to be provided. We can use the alias similar to a normal type, as shown in the usage example. Parametric aliases are quite useful to create reusable abstractions with some fixed parameters and some flexible parameters, similar to partially bound types. The difference between them is that with parametric aliases we can add additonal constraints on the parameters, while with partially bound types we cannot add additional constraints on the parameters. For example, in the StringKeyedDict alias we have added a constraint that the value type T must implement the traits Intable and Copyable. This means that we can only use the StringKeyedDict alias with value types that implement those traits, which is not possible to express with a partially bound type.

12.19. Dependent types

Many programming languages have a clear distinction between types and values. Types are used to define how data is structured, while values are the actual data that we manipulate in our programs. Mojo blurs the line between types and values by allowing types to depend on values. This is a powerful feature that allows us to express complex relationships between types and values, and to define invariants in the type system. In short, dependent types are types that depend on values. Mojo allows us to define dependent types using parameters.

# tag::dependent_types[]
def append[N: Int, M: Int](a: InlineArray[Int, N], b: InlineArray[Int, M]) -> InlineArray[Int, M + N]: # The return type has a dependent type that depends on the parameters N and M
    var result = InlineArray[Int, M + N](fill=0)
    for i in range(N):
        result[i] = a[i]
    for i in range(M):
        result[N + i] = b[i]
    return result
# end::dependent_types[]

def main():
    # tag::dependent_types_use[]
    var a: InlineArray[Int, 3] = [1, 2, 3]
    var b: InlineArray[Int, 2] = [4, 5]
    var c: InlineArray[Int, 5] = append(a, b)
    print(c) # [1, 2, 3, 4, 5]
    # var d: InlineArray[Int, 4] = append(a, b) # Uncommenting this line will cause a compile-time error 
    # end::dependent_types_use[]

Usage:

    var a: InlineArray[Int, 3] = [1, 2, 3]
    var b: InlineArray[Int, 2] = [4, 5]
    var c: InlineArray[Int, 5] = append(a, b)
    print(c) # [1, 2, 3, 4, 5]
    # var d: InlineArray[Int, 4] = append(a, b) # Uncommenting this line will cause a compile-time error

In the above example, we have a function append that takes two arrays with different lengths, and returns a new array that is the concatenation of the two input arrays. The return type of the function is a dependent type that depends on the lengths of the input arrays. The length of the returned array is the sum of the lengths of the input arrays. This allows us to define the output type in relation to the input types, and to ensure that the function is used correctly with respect to the lengths of the input arrays. If we try to use the append function with arrays of incompatible lengths, we would get a compile-time error, as the dependent type would not be satisfied. This helps us to catch errors at compile-time, which would have been caught at runtime in other programming languages without dependent types.

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

12. Parameters and compile-time programming

12.1. Compile-time execution of code

12.2. Conditional execution at compile-time

12.3. Compile-time if

12.4. Compile-time for

12.5. Compile-time expressions

12.6. Parameters in functions

12.7. Keyword parameters

12.8. Inferred-only parameters

12.9. Variadic parameters

12.10. Default values in parameters

12.11. Parameters in structs

12.12. Trait composition

12.13. Conditional conformance

12.14. where clauses

12.15. Custom compile-time checks

12.16. Parameterized errors

12.17. Partially bound parameters

12.18. Parametric aliases

12.19. Dependent types

12.3. Compile-time `if`

12.4. Compile-time `for`

12.14. `where` clauses