Comptime parameterized return type

Mojo
1

Use case:

def expect[ReturnToken: Bool = False](mut self, t: Token.Type)
    -> Token if ReturnToken else None:

    ref cur = self.current()

    if unlikely(cur.type != t):
        raise Error("Unexpected token type " + Token.type_name(cur.type) + 
                    " (expected " +  Token.type_name(t) + ")")

    self.skip()

    if ReturnToken:
        return cur.copy()

Reading https://docs.modular.com/mojo/manual/metaprogramming/, I couldn’t find an answer on how to specify a different return type depending on the ReturnToken parameter.

Token if ReturnToken else None does not work.

2

There is now a way in the latest nightlies to do this modular/mojo/stdlib/std/utils/type_functions.mojo at main · modular/modular · GitHub

1 Like
3

Thanks, I tried

def expect[ReturnToken: Bool = False](mut self, t: Token.Type)
    -> ConditionalType[If=ReturnToken, Then=Token, Else=NoneType]:

    ref cur = self.current()

    if unlikely(cur.type != t):
        raise Error("Unexpected token type " + Token.type_name(cur.type) + 
                    " (expected " +  Token.type_name(t) + ")")

    self.skip()

    if ReturnToken:
        return cur.copy()

but I still get:

parser.mojo:332:25: error: parameter 'Else' has 'AnyStruct[Token]' type, but value has type 'NoneType'
4

Hmmm seems like there may be some compiler limitations around using ConditionalType as a return type. This is the only way I could get this to work right now.

from std.utils.type_functions import ConditionalType


comptime Cond[b: Bool] = ConditionalType[
    Trait=ImplicitlyCopyable, If=b, Then=Int, Else=NoneType
]


def expect[ReturnToken: Bool = False](t: Int) -> Cond[ReturnToken]:
    comptime if ReturnToken:
        return rebind[Cond[ReturnToken]](t)
    else:
        return rebind[Cond[ReturnToken]](None)


def main():
    print(expect[True](1))  # 1
    print(expect[False](1))  # None
2 Likes
5

Thing I noticed: 13 people clicked on the link the last seven hours on Saturday…

1 Like
6

When @bgreni talks about the comptime-arts, we listen!

I thought I followed nightly stuff pretty well, but I totally missed ConditionalType.

2 Likes
7

Thanks Brian, that does seem to be the best possible option for now.

In my project, I’ve decided to switch back to 2 separate methods for now - switching to nightly results in compatibility issues with def raises and third party libraries (EmberJson). But it is good to know for future reference.

For the Mojo team: Would it be possible to enable the more intuitive syntax like I started with above?

def expect[ReturnToken: Bool = False](mut self, t: Token.Type)
→ Token if ReturnToken else None:

8

FYI: Add draft proposal: Comptime parameterized return type syntax by jbemmel · Pull Request #6178 · modular/modular · GitHub

1 Like
9

To vary the return type of a function you can just use a parameter for the actual return type.

Then inside the function assign to the parameterised return type using rebind so that the assignment of a concrete type to the parameterised return type works (the types are different names but boil down to the same type).

The following works in mojo, put it into test_return.mojo file:

from std.sys.intrinsics import _type_is_eq
from std.testing import TestSuite


def do_stuff[R: AnyType & ImplicitlyCopyable]() raises -> R:
    var result: R

    comptime
    if _type_is_eq[R, Int]():
        print("R is Int")
        var r = Int(42)
        result = rebind[R](r)
    elif _type_is_eq[R, Float64]():
        print("R is Float64")
        var r = Float64(3.14)
        result = rebind[R](r)
    elif _type_is_eq[R, NoneType]():
        print("R is NoneType")
        var r = None
        result = rebind[R](r)
    else:
        raise Error("Return type must be Int, Float64 or NoneType")

    return result


def test_stuff() raises:
    print("Int", do_stuff[Int]())
    print("Float64", do_stuff[Float64]())
    print("None", do_stuff[NoneType]())



def main() raises:
    TestSuite.discover_tests[__functions_in_module()]().run()

Test with mojo run test_return.mojo

HTH

1 Like
10

Interesting variation, thanks for sharing

11

Definitely doable, the only reason we need a ConditionalType ATM is that mojo compiler does not know how to merge two type to a common trait bound.

We might even be able to eliminate rebind[Cond[ReturnToken]] when the return is in a comptime if. Feel free to file a feature request :slight_smile:

12

See [Feature Request] [compiler] Support for Comptime-Parameterized Return Types · Issue #6188 · modular/modular · GitHub