Function Overloading in Mojo

functions
methods
overloading
Author

Kaizen Engineering

Published

January 26, 2024

Overview

Mojo allows multiple functions with the same name to coexist, provided their argument or parameter signatures are different. Function overloading features are seen in other languages as well, such as C++.

Function overloading allows semantically similar functions to share a name, facilitates easier customization of implementations based on argument or parameter signatures, and empowers the compiler to more rigorously type check at the function call site.

This article covers Mojo syntax for function overloading. In future articles we may dive into details of overload resolution at function call site and discuss best practices for overloading.

To begin, let’s examine overloading via argument signatures.

A note on code examples

The code examples in this article provide function signatures and invocations, and they typically skip implementation details. This is to maintain a focus on function overloading without distracting particulars of data structures and algorithms.

Argument signature overloading

Consider a struct representing a list of strings: 

struct StringList:
    # Create empty list
    fn __init__(inout self): ...

    # Overload 1
    # Insert `element` at end of list 
    fn insert(inout self, element: String):
        print("In Overload 1")
    # Overload 2
    # Insert `element` at `index`. Raises exception if `index` is outside [0, len],
    # where len is current length of this `StringList`
    fn insert(inout self, element: String, index: Int):
        print("In Overload 2")
    # Overload 3
    # Insert all `elements` at end of list
    fn insert(inout self, elements: StringList):        
        print("In Overload 3")
1
Overload 1
2
Overload 2
3
Overload 3

The three insert functions share the same name. They are legal overloads since their arguments differ by count or by type. Overload 1 and Overload 2 differ in argument counts, with two and three arguments respectively. And Overload 1 and Overload 3 differ in the type of second argument — String and StringList respectively. Based on the number and type of values at the call site, the compiler attempts to resolve the call to one of the overloads. Let’s examine this next.

Call resolution

The Mojo compiler resolves the call at invocation site, selecting one of the overloads based on the argument list. Let’s examine the output of executing the insert call in this code:

var tasks: StringList = StringList()
let task1: String = String("Review proposed design")

# Call Site A
tasks.insert(task1)
1
Call Site A 
In Overload 1

The output shows that Overload 1 was invoked.

At Call Site A, insert is invoked with two values. The first is the implicit value tasks of type StringList and the second is the explicit value task1 of type String. Based on the number of values, both Overload 1 and Overload 2 are candidates. Among these two, only Overload 1 matches the argument types. Consequently, the compiler resolved the call at Call Site A to Overload 1.

Let’s take a look at another example:

 # Call Site B
tasks.insert(task1, 0)
1
Call Site B 
In Overload 2

At Call Site B insert is invoked with three values, of types StringList, String, and IntLiteral. This matches Overload 2. Note that the compiler implicitly converts the literal 0 from type IntLiteral to Int. We’ll discuss type conversions for call resolution in greater depth later in a later article.

No overloading on result type

Mojo functions cannot be overloaded on result type alone. Let’s look at what happens when this struct declaration is compiled:

struct StringList:
    # Overload 1
    # Count all items
    fn count(self) -> UInt32: ...
    # Overload 2
    # Count all items
    fn count(self) -> UInt64: ...
1
Overload 1
2
Overload 2 
error: Expression [24]:8:5: redefinition of function 'count' cannot overload on return type only
    fn count(self) -> UInt64: ...   # <2>
    ^

Expression [24]:5:5: previous definition here
    fn count(self) -> UInt32: ...   # <1>
    ^

expression failed to parse (no further compiler diagnostics)

The compiler generated an error. The accompanying message indicates that Overload 2 of count is not legal since Overload 1 already exists.

Overloaded constructors

Mojo allows constructors to be overloaded as well. For example:

struct StringList:
    # Overload 1
    # Create empty list
    fn __init__(inout self):
        print("In Overload 1")
    # Overload 2
    # Create list initialized with `count` number of copies of `element`
    fn __init__(inout self, element: String, count: Int):
        print("In Overload 2")
    # Overload 3
    # Create list initialized with strings from `elements` between `start_index` and    
    # `end_index`
    fn __init__(inout self, elements: StringList, start_index: Int, end_index: Int):
        print("In Overload 3")
1
Overload 1
2
Overload 2
3
Overload 3

The argument lists of the three overloads differ by count or types, in keeping with Mojo’s overloading rules.

Here are invocations of each of them:

# Call Site A
let names = StringList()
# Call Site B
let statuses = StringList("Pending", 3)
# Call Site C
let copied_statuses = StringList(statuses, 1, 3)
1
Call Site A
2
Call Site B
3
Call Site C 
In Overload 1
In Overload 2
In Overload 3

As we can see, the compiler resolves Call Site A, Call Site B, and Call Site C respectively to Overload 1, Overload 2, and Overload 3.

Parameterized functions

When invoked with different types, a parameterized function effectively generates overloads with a shared function name.

Consider the simple parameterized swap function below that exchanges the referenced values:

# Swap the values of `value1` and `value2`
fn swap[valueType: Copyable](inout value1: valueType, inout value2: valueType):
    let temp = value1
    
    value1 = value2
    value2 = temp

The compiler instantiates a function based on the type of arguments. Here’s an example with Int values:

var i1 = 2
var i2 = 3

print("Before swap:", i1, i2)
# Call Site A
swap(i1, i2)
print("After swap:", i1, i2)
1
Call Site A 
Before swap: 2 3
After swap: 3 2


And now we use the same parameterized function with String:

var s1 = "Mojo"
var s2 = "Programming"

print("Before swap:", s1, s2)
# Call Site B
swap(s1, s2)
print("After swap:", s1, s2)
1
Call Site B 
Before swap: Mojo Programming
After swap: Programming Mojo
Explicit parameter specification

Note that in swap invocations at Call Site A and Call Site B we relied on the compiler to determine the valueType parameter. An alternative is to spell out the parameter as shown below:

var i1 = 2
var i2 = 3

print("Before swap:", i1, i2)
swap[Int](i1, i2)
print("After swap:", i1, i2)


var s1 = "Mojo"
var s2 = "Programming"

print("Before swap:", s1, s2)
swap[StringLiteral](s1, s2)
print("After swap:", s1, s2)
Before swap: 2 3
After swap: 3 2
Before swap: Mojo Programming
After swap: Programming Mojo

Parameter signature overloading

Mojo functions can also be overloaded on parameter signatures. The parameters of the overloads just need to differ in type or arity.

Take a look at this StringUtils struct with format functions overloaded via parameter signatures. They are a rudimentary version of string formatting functions found in other languages such as Python and C++.

Each overload takes the same logical arguments. It uses the format_str argument as a template, and replaces each occurrence of substring “{}” by the remaining arguments in order, converted to strings.

struct StringUtils:    
    # Each overload of `format` returns a copy of the `format_str` with an  occurrence 
    # of substring "{}" replaced by the remaining argument, converted to strings.
    # There should be exactly one occurrence of the substring. If not, `format` raises 
    # an `Error`.
    #
    # Overload 2 extends overloads 1 by taking an additional `replacement_field` 
    # parameter that enables customization of the field demarcator to other than the 
    # default "{}".

    # Overload 1
    @staticmethod
    fn format[
        param1Type: Stringable
    ](
        format_str: String, 
        arg1: param1Type
    ) raises -> String:
        print("In Overload 1")

        return ""   # Placeholder return

    # Overload 2
    @staticmethod
    fn format[
        replacement_field: StringLiteral, 
        param1Type: Stringable
    ](
        format_str: String, 
        arg1: param1Type
    ) raises -> String:
        print("In Overload 2")

        return ""   # Placeholder return
1
Overload 1
2
Overload 2

These format functions are valid overloads since their parameters differ in count. Overload 1 has one parameter - param1Type, while Overload 2 has two - replacement_field and param1Type.

Let’s try some invocations next.

Call resolution

If the replacement_field value parameter is not used, the Mojo compiler attempts to resolve the call to Overload 1. Let’s examine the output of executing calls to format.

Our first application of format generates a labelled username:

# Call Site A
let username = StringUtils.format("Username: {}", "dev#1")
1
Call Site A 
In Overload 1

The output shows that Overload 1 was invoked.

While both overloads are candidates, Overload 1 is the simpler match and the one the compiler picked. We’ll dive into the compiler’s resolution rules in greater detail in a future article.

Let’s examine another example - this time leveraging the second overload, which carries the replacement_field parameter. Consider a case in which the format_str argument has the “{}” substring embedded, however it is to be preserved; not replaced. In this example, we have a JSON object with an name/value pair where the value is the empty JSON object “{}” and needs to be preserved. So, we instead use “$$” as the replacement field:

# Call Site B
let empty_object = StringUtils.format["$$"]('{"$$": {}}', "postal_address")
1
Call Site B 
In Overload 2

The output shows that Overload 2 was invoked.

While both overloads are candidates, Overload 2 is the only one that accepts a StringLiteral value parameter and the one the compiler picked.

Explicit parameter specification

Note that Call Site B did not explicitly pass in String as a parameter and the Mojo compiler successfully inferred the type automatically. In contrast, here is invocation with two explicit parameters:

let empty_object = StringUtils.format["$$", String]('{"$$": {}}', "postal_address")
In Overload 2

Combined argument and parameter signatures overloading

Mojo allows function overloading via both argument and parameter signature overloading at the same time.

Let’s expand our StringUtils example from the previous section to introduce more overloads of the format function to support multiple replacement fields in format_str. While the overloads vary in argument and parameter signature, they all serve the purpose as before - format a string based on an input string template and replacements for demarcated fields:

struct StringUtils:    
    # Each overload of `format` returns a copy of the `format_str` with all occurrences 
    # of substring "{}" replaced by the remaining arguments in order, converted to 
    # strings. The number of occurrences of the substring must be the same as the 
    # number of remaining arguments. If not, `format` raises an `Error`.
    #
    # Overloads 2, 4, and 6 extend overloads 1, 3, and 5 respectively with an
    # additional `replacement_field` parameter that enables customization of the field
    # demarcator to other than the default "{}".

    # Overload 1
    @staticmethod
    fn format[param1Type: Stringable](
        format_str: String, 
        arg1: param1Type
    ) raises -> String:
        print("In Overload 1")

        return ""   # Placeholder return

    # Overload 2
    @staticmethod
    fn format[
        replacement_field: StringLiteral, 
        param1Type: Stringable
    ](
        format_str: String, 
        arg1: param1Type
    ) raises -> String:
        print("In Overload 2")

        return ""   # Placeholder return

    # Overload 3
    @staticmethod
    fn format[param1Type: Stringable, param2Type: Stringable](
        format_str: String, 
        arg1: param1Type, 
        arg2: param2Type
    ) raises -> String:
        print("In Overload 3")

        return ""   # Placeholder return

    # Overload 4
    @staticmethod
    fn format[
        replacement_field: StringLiteral, 
        param1Type: Stringable, 
        param2Type: Stringable
    ](
        format_str: String, 
        arg1: param1Type, 
        arg2: param2Type
    ) raises -> String:
        print("In Overload 4")

        return ""   # Placeholder return

    # Overload 5
    @staticmethod
    fn format[
        param1Type: Stringable, 
        param2Type: Stringable, 
        param3Type: Stringable
    ](
        format_str: String, 
        arg1: param1Type, 
        arg2: param2Type,
        arg3: param3Type
    ) raises -> String:
        print("In Overload 5")

        return ""   # Placeholder return

    # Overload 6
    @staticmethod
    fn format[
        replacement_field: StringLiteral, 
        param1Type: Stringable, 
        param2Type: Stringable, 
        param3Type: Stringable
    ](
        format_str: String, 
        arg1: param1Type, 
        arg2: param2Type,
        arg3: param3Type
    ) raises -> String:
        print("In Overload 6")

        return ""   # Placeholder return
1
Overload 1
2
Overload 2
3
Overload 3
4
Overload 4
5
Overload 5
6
Overload 6

Call resolution

Let’s take a look at the output of executing some calls to format. The following generates a JSON object with one name/value pair:

# Call Site C
let key_value = StringUtils.format("format {} {}", "usage_count", 6)
1
Call Site C 
In Overload 3

The output indicates that Overload 3 was invoked.

Call Site C invokes format with no ‘StringLiteral’ value parameter and three arguments. So the compiler selects Overload 3 was invoked.

As another example, we format a log entry with a standard format for date/time, log level, and log text.

# Call Site D
let log_entry = StringUtils.format(
    "[{}] {}: {}", 
    "2024-01-11T19:20:52Z", 
    3, 
    "File write completed."
)
1
Call Site D 
In Overload 5

The output demonstrates that Overload 5 was invoked.

Call Site D invokes format without replacement_field value parameter and four arguments. Since Overload 5 is the only candidate without a StringLiteral value parameter and 4 arguments, it is the only one examined further by the compiler. The first argument is a String and the remaining of Stringable types as required by that overload. Consequently, the compiler fully resolved the call to Overload 5.

Summary

Function overloading allows multiple functions to share a name. This enables meaningful grouping of functions based on their names, alternative implementations when needed, and stronger type checking at call sites.

Mojo supports overloading of functions via distinctive argument and parameter signatures. This capability extends to constructors as well. The compiler attempts to automatically resolve a call to an overloaded function based on types and arity of arguments and parameters.

We will examine Mojo’s overload resolution rules in greater depth in a future article.

Resources

Mojo

  1. Official page for the Mojo programming language.
  2. Official documentation on function overloading on arguments in Mojo.
  3. Official documentation on function overloading on parameters in Mojo.

Function overloading in C++

Note that in C++ the words argument and parameter have different and overlapping meanings when compared to their meanings in Mojo. In C++, variables defined in function declaration and used within function definition are called parameters, while values used for function invocation are called arguments.

  1. Overview of function overloading in C++ and coverage of overload resolution in C++.

Template functions in C++ are the closest feature to parameterized functions in Mojo.

  1. Overview of function templates in C++ and brief coverage of function template resolution in C++.

Function overloading in Python

Python does not natively feature overloaded functions and compile-time call resolution. However, approximations are available.

  1. Dispatch on type of first argument via functools.singledispatch.
  2. Dispatch on types of multiple arguments via multimethod or multipledispatch.

String formatting functions

These functions are available in many programming languages.

  1. String formatting in Python.
  2. String formatting in C++.