Mojo compile time programming - specifying a type

[pcupit]

I run the following and get an error, not sure what's going on here
error: Expression [35]:16:30: 'type' does not implement the 'lt' method return self.first < rhs.first or

struct MyPair[type:AnyType]:
var first: type
var second: type

# We use 'fn' instead of 'def' here - we'll explain that soon
fn __init__(inout self, first: type, second: type):
    self.first = first
    self.second = second


    
fn __lt__(self, rhs: MyPair[type]) -> Bool:       
       return self.first < rhs.first or
          (self.first == rhs.first and
           self.second < rhs.second)
    
fn __copyinit__(inout self, other: MyPair[type]):
    self.first = other.first
    self.second = other.second

def pairTest() -> Bool:
let p = MyPair[Int](1, 2)
let q = MyPair[Int](3, 4)

return p < q 

[DayDun]

type here could be any type. If you pass it a dictionary type for example, it doesn't really make sense to call self.first < rhs.first, that's why it's complaining. The proper solution here is traits, so that you can say type isn't just any type, it is any type that can be compared. But until traits are implemented you'll have to use DType or something to specify numeric types.

[Moosems]

When traits are implemented, how would I specify this kind of thing 😄?

[rd4com]

It works with compile time meta programming @pcupit @DayDun @Moosems :

see: #859 for other examples

struct MyPair[type:AnyType]:
    var first: type
    var second: type

    # We use 'fn' instead of 'def' here - we'll explain that soon
    fn __init__(inout self, first: type, second: type):
        self.first = first
        self.second = second

    fn __lt__(self, rhs: MyPair[type]) raises-> Bool:  
        @parameter
        if type == Int:
            return rebind[Int](self.first) < rebind[Int](rhs.first) or
              (rebind[Int](self.first) == rebind[Int](rhs.first) and
               rebind[Int](self.second) < rebind[Int](rhs.second))
        raise("not implemented")
        return False

               
    fn __copyinit__(inout self, other: MyPair[type]):
        self.first = other.first
        self.second = other.second

def pairTest() -> Bool:
    let p = MyPair[Int](1, 2)
    let q = MyPair[Int](3, 4)
    return p < q 
    
def main():
    print(pairTest())

[Moosems]

It would be nice if there was a type() function 😄.

[lattner]

Right, exactly!

[cheburakshu]

@DayDun / @pcupit Please consider revising the answer. @Moosems here is how I impemented the OP's question using generics. Mojo doesn't allow __copyinit__ method for register_passable but this works with Mojo hosted jupyter notebook's version. Full code below:

@register_passable("trivial")
struct MyPair[T: DType]:
    var first: SIMD[T, 1]
    var second: SIMD[T, 1]    

    @always_inline
    fn __init__(first: SIMD[T, 1], second: SIMD[T, 1]) -> Self:
        return MyPair[T] {first: first, second: second}

    @always_inline
    fn __lt__(self, rhs: MyPair[T]) -> Bool:       
        return self.first < rhs.first or
            (self.first == rhs.first and
             self.second < rhs.second)

    
let a = MyPair[DType.int32](1, 2)
let b = MyPair[DType.int32](3, 4)
print('Expect: True, Actual: ', a < b)

Prints:

Expect: True, Actual:  True

[Moosems]

Would you try making something similar that uses a Pointer and has the memory properly handled (especially in move init and copyinit)?

[cheburakshu]

@register_passable("trivial")
struct MyPair[T: DType]:
    var first: Pointer[SIMD[T, 1]]
    var second: Pointer[SIMD[T, 1]]   

    @always_inline
    fn __init__(owned first: SIMD[T, 1], owned second: SIMD[T, 1]) -> Self:
        return MyPair[T] {first: Pointer[SIMD[T, 1]].address_of(first), second: Pointer[SIMD[T, 1]].address_of(second)}

    @always_inline
    fn __lt__(self, rhs: MyPair[T]) -> Bool:       
        return self.first.load(0) < rhs.first.load(0) or
            (self.first.load(0) == rhs.first.load(0) and
             self.second.load(0) < rhs.second.load(0))


var a = MyPair[DType.int32](1, 2)
let b = MyPair[DType.int32](3, 4)
print('Before Copy: Expect: True, Actual: ', a < b)
a = b
print('After Copy : Expect: False, Actual: ', a < b)

Output:

Before Copy: Expect: True, Actual:  True
After Copy : Expect: False, Actual:  False

[Moosems]

I can't help but notice the difference in "After Copy", is that supposed to happen?

[cheburakshu]

ah, it is a typo, good catch. should beAfter Copy : Expect: False, Actual: False

[Moosems]

That's really cool! I wonder if I can use this logic in my Array type :).