[BUG]: Behaviour of the type checker is inconsistent (some expressions are manipulated prior to type-checking)

[nmsmith]

Bug description

I've recently been trying to understand how the type checker reasons about the equality of types. I've noticed a few inconsistencies and potential bugs. These are demonstrated below.

In short, I think the following things might be bugs:

1. Some MLIR instructions appear to be executed prior to type checking, which allows type expressions that are syntactically distinct to compare as equal. (See example f4_success below.) Furthermore, inlining appears to occur prior to type-checking. (See example f5_failure and f6_success.) I don't believe this should be happening. It would make more sense for the type checker to rely exclusively on symbolic reasoning, using the original, untransformed AST, because:
   • In the general case, the type checker cannot evaluate expressions, because those expressions may depend on function parameters.
   • The programmer (not just the compiler) needs to be able to reason about what programs will type-check successfully. The programmer is only able to reason in terms of the untransformed AST. Therefore, the compiler should probably do this as well.
2. Generally, the compiler seems to perform symbolic substitution of alias variables prior to testing for type equality. This transformation does make sense to me, because this allows people to use abbreviated type names without breaking type checking. More generally, it allows people to give a short name to a verbose parameter expression that needs to be used several times. However, substitution of aliases is not done in the following situations:
   • If a type alias is annotated with a trait bound, e.g. AnyType. (See example f2_failure below.)
   • If the alias is not a type. (See example f3_failure below.)

The following program demonstrates these potential bugs:

@value
struct Foo[x: AnyType]: pass

fn f1_success():
    """Type aliases are substituted prior to type checking."""
    alias A = Int
    var x = Foo[A]()
    var y = Foo[Int]()
    x = y

fn f2_failure():
    """However, annotating an alias prevents substitution."""
    alias A: AnyType = Int
    var x = Foo[A]()
    var y = Foo[Int]()
    x = y  # Error: cannot implicitly convert 'Foo[Int]' value to 'Foo[A]' in assignment

@value
struct Bar[x: Int]: pass

fn f3_failure():
    """If an alias isn't a type, it isn't substituted."""
    alias a = Int(4)
    var x = Bar[a]()
    var y = Bar[Int(4)]()
    x = y  # Error: cannot implicitly convert 'Bar[4]' value to 'Bar[a]' in assignment

fn f4_success():
    """Some expressions appear to be evaluated before type-checking."""
    var x = Bar[1 + 3]()
    var y = Bar[2 + 2]()
    x = y

fn f5_failure():
    """Ordinary function calls are NOT evaluated before type-checking."""
    var x = Bar[add(1, 3)]()
    var y = Bar[add(2, 2)]()
    x = y  # Error: cannot implicitly convert 'Bar[add(2, 2)]' value to 'Bar[add(1, 3)]' in assignment

fn add(x: Int, y: Int) -> Int:
    return x + y

fn f6_success():
    """This can be 'fixed' by inlining the function. I suspect what is happening is that inlining
    exposes an MLIR `add` instruction whose operands are literals, and this instruction is
    executed (i.e. the AST is simplified) prior to type-checking.
    """
    var x = Bar[inlined_add(1, 3)]()
    var y = Bar[inlined_add(2, 2)]()
    x = y

@always_inline
fn inlined_add(x: Int, y: Int) -> Int:
    return x + y  # This call to `Int._add__` is inlined too.

fn note_1():
    """Expressions that aren't evaluated are compared symbolically."""
    var x = Bar[add(1, 3)]()
    var y = Bar[add(1, 3)]()
    x = y

fn note_2[k: Int]():
    """Symbolic reasoning is necessary, because expressions may depend upon function parameters."""
    var x = Bar[k + 3]()
    var y = Bar[k + 3]()
    x = y

Steps to reproduce

See above.

System information

Mojo 0.7.0 (af002202)

[laszlokindrat]

Thanks for opening the ticket! I don't see any of the behavior you described as bugs per se; AFAICT all these examples work as currently intended. I do agree that a lot of this is not necessarily intuitive, and perhaps @Mogball can provide some details on why this is the current behavior (e.g. why @always_inline functions are substituted early, and aliases aren't). Beyond that, I think this might be something we want to highlight in documentation. CC: @arthurevans

[Mogball]

A lot of this is work in progress, which includes some hacks added to make "common cases" work while we spin on general solutions (if they exist). I don't think we should necessarily document the hacks.

[soraros]

which includes some hacks added to make "common cases" work while we spin on general solutions (if they exist).

That's pretty understandable. Can we document roughly what suppose to work/not work? Maybe in the form of a design note instead of a section in the manual. Currently, there is no solid ground to build intuition upon.

[Mogball]

What @nmsmith "discovered" is actually pretty accurate, with the exception that alias values are only substituted when it is not declared in a function. This divergence in behaviour is "just" tech debt. The reason f1 type checks is because we can compare the metatypes of Int. I'm not sure why f2 fails.

The fact that inlined functions are evaluated before type checking is intentional, because it allows Foo[1+2] and Foo[3] to be equal, which is pretty important in lots of code, but we don't have any symbolic reasoning, so Foo[a+1-1] will not type check to be equal to Foo[a], which is kind of ridiculous. There was some discussion around symbolic simplifications of integers only, but that only goes so far.

[nmsmith]

Beyond special-casing the + function, I'm not sure that it even makes sense for functions to appear in types. At least, not in a programming language that has untracked side-effects.

Consider:

fn some_function(x: Int) -> Int: ...

fn foo[x: Int]():
    var s: SIMD[DType.bool, some_function(x)]
    s = SIMD[DType.bool, some_function(x)](0, 0, 0, 0)

This passes the Mojo 0.7 type checker. However, there's no guarantee that some_function() returns the same result every time it is called. If it doesn't, then the above program is not type-safe.

Mojo currently appears to "solve" this problem by disallowing side-effects during specialization, but I assume this constraint will be lifted at some point, because side-effects such as reading from a file will be very useful.

Type-safety of the above program can be restored by only invoking the function once, and using its value several times:

fn some_function(x: Int) -> Int: ...

fn foo[x: Int]():
    alias y = some_function(x)
    var s: SIMD[DType.bool, y]
    s = SIMD[DType.bool, y](0, 0, 0, 0)

This suggests that perhaps only variables and literals should be allowed in types, because they are guaranteed not to have side-effects?

Update: I suppose we also could allow functions to appear in parameter lists if they are only ever invoked once:

fn foo[x: Int]():
    # This is fine, because `some_function` is only invoked once:
    var s = SIMD[DType.bool, some_function(x)](0, 0, 0, 0)
    # But we can't reassign `s` using an expression that invokes the function a second time:
    s = SIMD[DType.bool, some_function(x)](1, 1, 1, 1)
    # However, we can reassign it using the OLD parameter value:
    s = SIMD[DType.bool, s.size](1, 1, 1, 1)

But on top of this, I suppose it could be beneficial (or necessary) to imbue the compiler with the knowledge that + and - on the "standard" integer types are pure functions (as it currently assumes), and therefore two occurrences of the expression 4+1 (for example) are equal. That's the only way for the compiler to reason that two separately-manipulated arrays have the same length:

fn foo():
    var a = Array[Int, 4](...)
    var b = Array[Int, 4](...)
    var a2 = a.append(0)
    var b2 = b.append(0)

This doesn't require 4+1 to be evaluated at compile-time though. Syntactic reasoning is sufficient. I'd be keen to see some examples of where evaluation of + prior to type-checking is necessary. I don't know of any.

[Mogball]

Mojo currently appears to "solve" this problem by disallowing side-effects during specialization, but I assume this constraint will be lifted at some point, because side-effects such as reading from a file will be very useful.

Reading from a file is not really a side effect, and yeah we do intend to support that in a way that is sane and possible to reason about. The interpreter's side effect model is intentionally designed to ensure compile time code only does sane things. Observative mutability will probably never be allowed in comptime function calls.

This suggests that perhaps only variables and literals should be allowed in types, because they are guaranteed not to have side-effects?

1+2 is a function call, and most things in Mojo are function calls so we have to model this. We can model pure functions in the type system, for instance

But @Mogball says that evaluation of + prior to type-checking is necessary, and it is obviously a pure function, so maybe this can be a privileged function, that is allowed to appear in types. (I'd be interested in seeing some examples of where evaluation of + prior to type-checking is necessary. I don't know of any.)

Concating a tensor is an example of where you want the result type to be refined.

[nmsmith]

1+2 is a function call

@Mogball Yes, sorry, I amended my post with a bit more discussion about + shortly before you replied. I appreciate that it's a function call, not a built-in operation.

We can model pure functions in the type system, for instance

I'd love to see that, but that would obviously present an extraordinary challenge. (Last year, I wrote a massive post with half-baked ideas about purity on the discussion forum. That helped me realize how thorny an issue it is.)

[nmsmith]

Actually, I suppose even the use of literals in types can trigger side-effects, because passing a literal as a parameter often implies a conversion to another type:

struct Foo[x: Int32]: pass
# The following annotation:
var x: Foo[0]
# Is the same as:
var x: Foo[Int32(0)]

This is perhaps incontrovertible evidence that the compiler needs to reason about purity somehow.