Compile time check on __getitem__()

[mojodojodev]

Hi I was experimenting with compile time checking shapes using NDBuffer:

struct Tensor[rank: Int, shape: DimList, type: DType]:
    var data: DTypePointer[type]
    var buffer: NDBuffer[rank, shape, type]
    fn __init__(inout self):
        let size = shape.product[rank]().get()
        self.data = DTypePointer[type].alloc(size)
        memset_zero(self.data, size)
        self.buffer = NDBuffer[rank, shape, type](self.data)

    fn get[*idx: Int](self) -> SIMD[type, 1]:
        @parameter
        fn check_dim[i: Int]():
            assert_param[idx[i] < shape.value[i].get()]()

        unroll[rank, check_dim]()

        return self.buffer.simd_load[1](VariadicList[Int](idx))

So that something like:

let x = Tensor[3, DimList(2, 2, 2), DType.uint8]()
x.data.simd_store(0, SIMD[DType.uint8, 8](1, 2, 3, 4, 5, 6, 7, 8))
print(x.get[2, 0, 0]())

Throws a compile time error

Wondering if you've discussed potentially allowing this behavior by specifying parameters on __getitem__:

fn __getitem__[*idx: Int](self) -> SIMD[type, 1]:
    @parameter
    fn check_dim[i: Int]():
        assert_param[idx[i] < shape.value[i].get()]()

    unroll[rank, check_dim]()

    return self.buffer.simd_load[1](VariadicList[Int](idx))

Compile time shape checks inside [] would be really nice, and fit nicely with Mojo how everything in [] is a compile time parameter.

Even better would be the ability to overload it:

__getitem__(borrowed self, *idx: Int)

So you can do:

x(2, 0, 0) for no compile time checks if using dynamic values
x[2, 0, 0] for compile time checks