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

3. Functions

In the previous chapter you saw how a function print is being called in a program. In order to call a function, the function must be defined somewhere. In Mojo, there are two ways to define a function. The first way would be very familiar to Python developers. The second one is more closer to another language, Rust.

3.1. def

One way to define a function is by using the keyword def.

def my_function(text):
    print(text)

def main(): my_function("Hello World!")

def function definitions have the least ceremony. For beginners it is the easiest way to define functions and closely resembles how Python functions are defined. In the above code, def my_function(text): defined a function named my_function and declared that it takes an argument text. The following statement my_function("Hello World!") calls that function my_function and passes "Hello World!" as the argument where text was declared.

You may have noticed that after def my_function(text):, there is a new line and 4 spaces before the statement print(text). This is because Mojo, like Python, uses whitespace indentation to demarcate the function’s body. Many mainstream programming languages use braces "{}" for function body. However, Mojo uses indented whitespace and is particular about the whitespace being aligned. You will find this syntax not only for function bodies, but also for other statements that expect a block. We will come to those cases in later chapters. Please notice that the statement my_function("Hello World!") appears differently aligned than print(text). This is because my_function("Hello World!") is not part of the `my_function’s body itself.

The following illustration shows the simplified structure of a Mojo function.

3.2. fn

The other way to define is function is by using the keyword fn.

fn my_function(text: StringLiteral):
    print(text)

fn main(): my_function("Hello World!")

While the def based function definitions demands only the most essential elements, the fn based function definitions are more strict in nature. The reason for this is simple. The fn functions are intended to be high performance and for it to be high performant, it needs to provide the Mojo compiler with much more details. The details help the Mojo compiler to create an optimized version of the executable file.

What are those details that help Mojo compiler? The main element that helps Mojo compiler is something called the "type" of a variable. In the above example, you see the definition text: StringLiteral. In comparison to the def version of the function, we have a new declaration : StringLiteral. This is called a type annotation. This declaration says that the function argument text will contain only String literals. More importantly, it will never contain any other content than String literals. This gives the Mojo compiler a very important hint. Without such a hint, the Mojo compiler has to accommodate many different scenarios. For example, if we do not declare that the text is of type StringLiteral, it will have to assume that the text may contain numbers or other types of objects. Then it has to generate a very generic executable code that is able to handle many other types of values. However, when we tell Mojo compiler that the text will take only StringLiterals, it can generate a very specific and highly optimized code that handles only StringLiterals.

The body of the fn functions are also demarcated by whitespace similar to def functions.

The following illustration shows the simplified structure of a fn function.

You may have noticed one additional concept that we did not yet touch upon. The return type. Both fn and def functions support returning values from the function. So far we have not used it. However, in many cases we want to call functions to perform calculations and after the calculations are completed, we expect a result as output from the function. The following example shows such a use case.

def add(a, b):
    return a + b

def main(): print(add(1, 2))

If you run the above program, you get the output as 3.

The function add took two arguments and performed an addition. Then it returned that result with keyword return, which was then shown on the terminal using the print function. The statement print(add(1, 2)) actually made two function calls. The inner function call add got executed first and then the outer function call print. This is called the nesting of function calls. The result from the inner function call was passed to the outer function call.

mojo addition_return_def.mojo
3

Now, the same using fn function definition.

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

fn main(): print(add(1, 2))

If you run the above program, you get the output as 3. Here you may have noticed that the type of the values are being explicitly declared. The arguments passed to the function are both Int types and the return value of the function is also of type Int. The Int type represents all integer values, including both positive and negative values. The actual value is passed during the function call statement (call site).

mojo addition_return_fn.mojo
3

In case of fn function, we indicated in the function definition that the function will return a value of type Int.

3.3. Benefits of `fn` over `def`

As mentioned earlier, one of the benefits of using fn is to provide information to the compiler that would be used by the compiler to produce a highly performance code. The second benefit is to enforce program correctness.

Let’s look into the following program.

def add(a, b):
    return a + b

def main(): print(add(1, 1.5))

When you execute the above program, you would see the result as 2.5. Behind the scenes, the expression add(1, 1.5) added two different types. The first argument was of type Int and the second argument was of type Float. The addition’s result is of type Float.

FloatLiteral vs Float

I am glossing over when I say that 1.5 is a Float. Since it is provided in the source code directly, its actual type is FloatLiteral.

The above program worked fine for us. But what if we wanted to restrict addition to only integers? How can we prevent someone from passing Float values to add function?

The answer is by using fn.

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

fn main(): print(add(1, 2.5))

If you now execute the above program, you would get a compiler error (you may have slightly different error message depending upon the version of the Mojo compiler used).

addition_return_fn_2.mojo:4:21: error: invalid call to 'add': argument #1 cannot be converted from 'FloatLiteral' to 'Int'

The Mojo compiler complains because we have strictly restricted the argument type of b to Int and Mojo compiler cannot convert the passed value to Int. You have a guarantee from the Mojo compiler that only those values that can be converted safely into Int will be allowed to be passed. We say that the Mojo compiler ensures type safety when we use fn.

3.4. Benefits of `def` over `fn`

If fn is safer than def, then why should we use def at all? def has its uses. For example, if you want to prototype something and you are not sure what types to use, you can leave that decision for a later time and focus on the algorithm itself. Sometimes you really need the flexibility and dynamism of Python. In this case, def is the most appropriate way to go.

Mojo treats both def and fn styles as first class and both of them are useful in different contexts.

3.5. Mixing and matching

Although def does not require you to provide types, it does not prevent you from declaring types. For example, the following will have the same effect as the above mentioned fn version.

You may wonder what is the type of a in the def function. Its type is object. All values that do not have an explicit type is by default assigned as object type in Mojo. We will come back to it in detail later on in subsequent chapters.

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

def main(): print(add(1, 1.5))

error: invalid call to 'add': argument #1 cannot be converted from 'FloatLiteral' to 'Int'

3.6. Default return types

In the def function, when you omit the return type annotations then Mojo assumes that the return type is object. However, in fn functions, when you omit the return type, then Mojo defaults to None, which indicates that there is no value being returned.

3.7. `main` function

There is a special function called main in Mojo, which is the function that is used by the compiler to determine the main entrance of a program. When you call a program executable created by Mojo, the very first function that is executed is the main function.

def main():
    print("Hello World!")

If you execute the above program, the main function will be called automatically and the text Hello World! will be printed.

3.8. Different styles of writing functions

A function can be defined in different styles.

def func1(r): ...

def func2(): pass

def func3(): print("Hello World!"); print("Good bye!")

def func4(): 
    pass

def main(): ...

Note that all the above style is also valid for fn functions. fn functions are the same as def functions, except for the difference that it demands type annotations. You can define the function body in the same line as the function definition only if the whole body is just a single line.

3.8.1. Semicolon

In the source code shown above, you may have noticed a semicolon in the body of function func3. Semicolon can be used to separate statements, which allow the statements to be written in a single line. Mojo follows the philosophy of Python, so use semicolon sparingly and only when it improves reading and understanding of code.

3.8.2. Ellipses in functions

In Mojo you would often see … defined in the function body. Ellipsis is just a built-in constant in Mojo, and is a placeholder. Within a function, it just means that the body is not yet implemented and the Mojo compiler will not complain about the missing body. Ellipses have other uses and we will cover them in subsequent chapters.

3.8.3. `pass` in functions

The pass keyword has a similar role as … in functions. It particularly tells the compiler that the implementation has been omitted.

One good rule of thumb is to use pass where you know that there is no need for an implementation and use … when you are expecting some implementation in the future (or in inherited entities - we will come to that later).

3.9. Arguments passed to functions

Arguments passed to a function cannot be modified within the function. Such arguments are said to be immutable as the function body cannot modify their value. This kind of restriction is helpful in large programs as the code that calls the function does not get surprised that the value it passed to a function has suddenly changed unexpectedly.

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

fn main(): print(add(1, 2))

Executing the above listed code results in compilation error.

error: expression must be mutable in assignment
    a = a + 1
    ^
mojo: error: failed to parse the provided Mojo

3.10. Nested functions

Mojo supports nesting of function within other functions. This applies to both fn and def style functions.

Nesting of functions limit the scope of the nested function to just the outer function. For example in the following code listing, it is not possible to call nested within main.

fn outer():
    fn nested():
        print("I am nested")
    nested()

fn main():
    outer()

3.11. Keyword arguments

Sometimes when functions take lots of arguments, it is much more clear if the name of the argument is specified when calling the function. Mojo supports keyword arguments, which is basically the ability to specify the argument name when we assign a value to that argument during a function call.

fn my_function(first: Int, second: Int) -> Int:
    return first + second

fn main():
    print(my_function(first = 1, second = 2))
    print(my_function(second = 2, first = 3))
    print(my_function(4, second = 5))

In the previous listing, we can see there are three different ways to call the function with keyword arguments. The first one specifies the name of the both arguments when passing the value. The second call demonstrates that when using the keyword arguments, the order of the arguments does not matter, as Mojo knows with the name itself which argument gets which value. In the third call, we see that we can mix and match positional argument with the keyword argument; however, here the order is important as the positional values must appear in the order in which they were declared in the function definition. Keyword arguments follow positional arguments.

What if you as an API designer want some arguments to be always specified positionally? In this case, you can enforce positional arguments by using a special argument, /.

fn my_function2(first: Int, second: Int, /) -> Int:
    return first + second

fn main():
    print(my_function2(first=1, second=2)) # compiler error

Executing the above listed code results in compilation error.

error: invalid call to 'my_function2': positional-only arguments passed as keyword operands: 'first', 'second'

You can also mix and match position-only arguments with keyword arguments. The / can be the last argument, in which case all the function arguments would be position-only. It cannot be the first argument though.

fn my_function(first: Int, second: Int, /, third: Int) -> Int:
    return first + second + third

fn main():
    print(my_function(1, 2, third=3))
    print(my_function(1, 2, 3))

Here the first two arguments are strictly position-only, while the third can be passed as keyword or positional as desired.

What if you wanted some arguments to be always keyword only? In this case, you can enforce keyword arguments by using the special argument, *.

fn my_function(first: Int, *, second: Int, third: Int) -> Int:
    return first + second + third

fn main():
    print(my_function(1, second=2, third=3))
    #print(my_function(1, 2, 3)) # Uncommenting would result in compiler error.

Similar to /, you can also mix and match keyword-only arguments with positional arguments. The * can be the first argument, in which case all the function arguments would be keyword-only. It cannot be the last argument though.

Keyword arguments make APIs ergonomic, as the programmer does not have to remember in which position what value need to be passed. It improves code readability and maintainability. It also reduces accidental mistakes when programmer wrongly assumes the order of the arguments.

Keyword arguments are applicable for both def and fn forms.

3.12. Default value

Mojo allows assigning default values to function arguments, which means when the caller does not pass a value to the argument, the function will take the given default value. That function argument therefore becomes optional for the caller. The default value must be of the same data type as the declared data type of the argument.

This feature is quite useful when defining ergonomic APIs, providing sensible default values for the caller, making the function easier to use.

fn deft_function(first: Int, second: Int = 10) -> Int:
    return first + second

fn main():
    print(deft_function(1)) # 'second' defaults to '10'
    print(deft_function(1, 2))

Default values are applicable for both def and fn forms.