Calculations in type parameters without pure

Suppose I have a type parameter M, which is a deterministic function of another parameter N. I would like to save it in the type, but not sure how to make this type stable without @pure. Is that even possible? MWE:

@inline calcM(x) = 2^x

struct MyType{N,M} end

function MyType{N}() where N
    @assert N β‰₯ 0
    M = calcM(N)
    MyType{N,M}()
end

@code_warntype MyType{3}() # not inferred of course
2 Likes

I don’t have an answer to this, but this is definitely not an of course.
If you just simplify the function to e.g.

@inline calcM(x) = x^2

then compiler decides to evaluate it during compile time:

julia> @code_warntype MyType{3}()
Variables
  #self#::Type{MyType{3, M} where M}
  M::Int64

Body::MyType{3, 9}
1 ─      Core.NewvarNode(:(M))
β”‚   %2 = ($(Expr(:static_parameter, 1)) β‰₯ 0)::Core.Const(true)
β”‚        %2
└──      goto #3
2 ─      Core.Const(:(Base.AssertionError("N β‰₯ 0")))
└──      Core.Const(:(Base.throw(%5)))
3 β”„      (M = Main.calcM($(Expr(:static_parameter, 1))))
β”‚   %8 = Core.apply_type(Main.MyType, $(Expr(:static_parameter, 1)), M::Core.Const(9))::Core.Const(MyType{3, 9})
β”‚   %9 = (%8)()::Core.Const(MyType{3, 9}())
└──      return %9

Along those lines: You can nudge the compiler to make your example type stable by adding an additional type check.

julia> function MyType{N}() where N
           @assert N β‰₯ 0
           T = typeof(N)
           M = calcM(N)::T
           MyType{N,M}()
       end

julia> @code_warntype MyType{3}()
Variables
  #self#::Type{MyType{3, M} where M}
  M::Int64
  T::Type{Int64}

Body::MyType{3, 9}
1 ─       Core.NewvarNode(:(M))
β”‚         Core.NewvarNode(:(T))
β”‚   %3  = ($(Expr(:static_parameter, 1)) β‰₯ 0)::Core.Const(true)
β”‚         %3
└──       goto #3
2 ─       Core.Const(:(Base.AssertionError("N β‰₯ 0")))
└──       Core.Const(:(Base.throw(%6)))
3 β”„       (T = Main.typeof($(Expr(:static_parameter, 1))))
β”‚   %9  = Main.calcM($(Expr(:static_parameter, 1)))::Core.Const(9)
β”‚         (M = Core.typeassert(%9, T::Core.Const(Int64)))
β”‚   %11 = Core.apply_type(Main.MyType, $(Expr(:static_parameter, 1)), M::Core.Const(9))::Core.Const(MyType{3, 9})
β”‚   %12 = (%11)()::Core.Const(MyType{3, 9}())
└──       return %12
1 Like

Thanks,

@inline calcM(x) = 1 << x

also works.

I am just not sure what the rules are. Is it because 2^x goes to power_by_squaring, which is recursive?

I could be wrong here, but isn’t this a good use case for generated functions?

@generated function MyType{N}() where N
  M = calcM(N)
  return quote
    MyType{N,$M}()
  end
end

seems to work okay with inference.

3 Likes