Scala 3 Macros: Possibly missing information in the AST for inlined lambda functions

[davesmith00000]

Hi! I don't know if this is really an issue or just my lack of understanding. Any insights welcome! Thank you.

Compiler version

3.5.0, also tested on 3.5.2.

Minimized example

Kindly provided by someone on Discord, the original version is at the bottom of the issue.

  // munit test
  test("wiring issue".only) {
    case class C(x: Float, y: Float)
    case class D(x: Float, c: C)

    inline def f: Float => C =
      foo => C(0f, foo)

    class Test:
      inline def g =
        List(3.14f).map: env =>
          val y = 2.0f
          D(env, f(y))

    println(DebugAST.anyToAST(Test().g))

    assert(1 == 2)
  }

Output

Inlined(
  None,
  Nil,
  Inlined(
    Some(Select(Apply(Select(New(TypeIdent("Test")), "<init>"), Nil), "g")),
    Nil,
    Typed(
      Apply(
        TypeApply(
          Select(
            Apply(
              TypeApply(Select(Ident("List"), "apply"), List(Inferred())),
              List(Typed(Repeated(List(Literal(FloatConstant(3.14f))), Inferred()), Inferred()))
            ),
            "map"
          ),
          List(Inferred())
        ),
        List(
          Block(
            List(
              DefDef(
                "$anonfun",
                List(TermParamClause(List(ValDef("env", Inferred(), None)))),
                Inferred(),
                Some(
                  Block(
                    List(ValDef("y", Inferred(), Some(Literal(FloatConstant(2.0f))))),
                    Apply(
                      Select(Ident("D"), "apply"),
                      List(
                        Ident("env"),
                        Inlined(
                          Some(Ident("f")),
                          Nil,
                          Block(
                            Nil,
                            Apply(Select(Ident("C"), "apply"), List(Literal(FloatConstant(0.0f)), Ident("foo")))
                          )
                        )
                      )
                    )
                  )
                )
              )
            ),
            Closure(Ident("$anonfun"), None)
          )
        )
      ),
      Inferred()
    )
  )
)

Expectation

In the original code, y is passed as an argument to f(..). f is a lambda with a named argument foo. f gets inlined, and continues to refer to foo in the AST, but there is no clue as to where foo came from, or that y (which is correctly defined) should be substituted in it's place.

Does that seem right? Perhaps it is use-case dependant. In my case I'm writing a transpiler, and I need to know where values come from.

More information

DebugAST is just doing Printer.TreeStructure.show(expr.asTerm) behind the scenes to give me the tree. Perhaps there's a better tool?

Original question posted on Discord

I have this code:

inline def fn1(v: Float): vec2 =
  vec2(v)

inline def fn2: Float => vec2 =
  alpha => vec2(0.0f, alpha)

inline def fragment: Shader[FragEnv, vec4] =
  Shader { env =>
    val y = 2.0f
    vec4(fn1(1.0f), fn2(y))
  }

And it produces this AST (just the relevant bit, starting with env =>):

DefDef(
  "$anonfun",
  List(TermParamClause(List(ValDef("env", Inferred(), None)))),
  Inferred(),
  Some(
    Block(
      List(ValDef("y", Inferred(), Some(Literal(FloatConstant(2.0f))))),
      Apply(
        Select(Ident("vec4"), "apply"),
        List(
          Inlined(
            Some(Apply(Ident("fn1"), List(Literal(FloatConstant(1.0f))))),
            Nil,
            Typed(
              Apply(
                Select(Ident("vec2"), "apply"),
                List(Inlined(None, Nil, Literal(FloatConstant(1.0f))))
              ),
              TypeIdent("vec2")
            )
          ),
          Inlined(
            Some(Ident("fn2")),
            Nil,
            Block(
              Nil,
              Apply(
                Select(Ident("vec2"), "apply"),
                List(Literal(FloatConstant(0.0f)), Ident("alpha"))
              )
            )
          )
        )
      )
    )
  )
)

fn1 is fine. It inlines the value and ignores the argument, but that isn't a problem.
fn2 ...talks about alpha, but the argument and type have vanished, and it's being called with y, but it hasn't inlined the value. So it appears broken / not wired up correctly in the AST.

[som-snytt]

Maybe it's just a "printer" problem, as it looks more reasonable under

scala-cli compile --server=false -S 3.5.2 -Vprint:refchecks -Yplain-printer inline-func.scala

or wait it only looks reasonable. There is a dangling _$1.

                    DefDef($anonfun,
                      List(List(ValDef(env, TypeTree(), Thicket(List())))),
                      TypeTree(),
                      Block(List(ValDef(y, TypeTree(), Literal(2.0f))),
                        Apply(Select(Ident(D), apply), List(Ident(env),
                          Inlined(Ident(f), List(),
                            Apply(Select(Ident(C), apply), List(Literal(0.0f),
                              Ident(_$1)))
                          )
                        ))
                      )

I can't see the Scala 3 forest for the trees. Or is it a thicket?

[davesmith00000]

Interesting! I didn't know you could do that. 💪

What's cool about that is that you do actually get the definition of f to look at - not sure it's fully expanded even here but maybe there's more there we could sub-print to get some of the types back, for instance.

However, fun as that is, if I understand what's happening here(?), we're getting the entire AST tree by running that command, but in my example I only get to see the sub-section relevant to my macro in which the definition of f is missing (I think).

[jchyb]

Running the example with scala-cli compile --server=false -S 3.5.2 -Xprint:inlining -Yplain-printer -Xprint-types allows us to also see the types of the trees, which for the relevant inlined fragment show us:

DefDef($anonfun,
  List(List(
    ValDef(env, TypeTree( | scala.this.Float),
      Thicket(List() | <notype>) | (env : scala.this.Float))
  )),
  TypeTree( | <empty>.this.D),
    Block(
      List(
        ValDef(y, TypeTree( | scala.this.Float),
          Literal(2.0f | (2.0f : scala.this.Float)) |
          (y : scala.this.Float))
      ),
      Apply(
        Select(Ident(D | <empty>.this.D.type), apply |
          (<empty>.this.D.apply :
            (x: <root>.this.scala.Float, c: <empty>.this.C):
              <empty>.this.D
          )
        ),
      List(Ident(env | (env : scala.this.Float)),
        Inlined(
          Ident(f |
            (<empty>.this.bbb$package.f :
              =>
                Function1[<root>.this.scala.Float,
                  <empty>.this.C]
            )
          ),
        List(),
          Block(List(),
            Apply(
              Select(Ident(C | <empty>.this.C.type), apply |
                (<empty>.this.C.apply :
                  (x: <root>.this.scala.Float, y:
                    <root>.this.scala.Float): <empty>.this.C
                )
              ),
            List(Literal(0.0f | (0.0f : scala.this.Float)),
              Ident(foo | (y : scala.this.Float))) |
              <empty>.this.C)
           | <empty>.this.C)
         | <empty>.this.C)
      ) | <empty>.this.D)
     | <empty>.this.D)
   | ($anonfun : (env: scala.this.Float): <empty>.this.D))

In the Ident(foo | (y : scala.this.Float))) part we see that inlined foo is typed with what seems to be a TermRef to y (which also should explain why we cannot see foo in the code/pretty printer).
I didn't test it yet but I imagine calling tpe on that Ident should return that TermRef.
Regardless, this still seems incorrect to me (although this might be an implementation detail, e.g. I know that inline vals are inlined via just setting a literal constant type and letting the optimisations do the rest of the work). I will be investigating further.

[davesmith00000]

Thank you for looking into it @jchyb. I shall check out the tpe and see if that yields anything in the meantime!

I've mostly just worked out how to do - erm... whatever I can do with scala 3 macros - by trial and error. So I'm confident there are lots of branches I've missed, don't know about, or just don't understand yet.

For example, I didn't know about the scala-cli command you've both used there. That seems to work nicely on standalone code, but my next question would be: Is there a way to run that one a piece of code in a larger code base? Or is the pretty printer I'm using the next best thing, or is there something better?

Thanks again!

[davesmith00000]

I know that inline vals are inlined via just setting a literal constant type

It does feel like this (well, this sort of thing) is where most of my problems stem from. I'm trying to use inline macros to write a transpiler, and instead of being able to work with the AST that represent the code that is there, I'm dealing with an AST that has already been partially inlined/optimised, and I've lost information.

Do Scala 3 macro annotations give you a non-inlined tree, do you happen to know? Or is it all hanging off the same underlying mechanism? Is it worth the hassle of moving from inline macros to annotation based macros, or even a compiler plugin?

[davesmith00000]

Answering my own quesiton: It looks to me like macro annotations are run before inlining. So when attempting to write a transpiler, this seems to leave me with 3 options with their pros and cons:

1. Full compiler plugin (a la Scala.js) - transform a whole program. I guess code sharing will need to happen by importing libs. The approach means I can't just write a shader within my game project (as we do now), it needs to be an external module / project that would emit GLSL code during compilation.
2. Macro annotations - Nice clean AST and you get to chose your scope level. However, code sharing seems impossible. You can transform everything under the current annotation, but you can't see anything outside of it. (Right?)
3. Inline macros - thanks to inlining, it's possible to use shared code, it's just clunky and the AST is garbled and you have to do a lot of creative interpretation.

Macro annotations would be the best solution, but being able to reuse code is kinda the whole point of the library (because code reuse in WebGL / GLSL 300 is difficult), and if we can't share code that I think I'm left with the current option, because I don't want to transform whole programs.

For the avoidance of doubt, here's what you get if you print the tree for this:

    case class C(x: Float, y: Float)
    case class D(x: Float, c: C)

    inline def f: Float => C =
      foo => C(0f, foo)

    @DebugAST2
    class Test:

      def g =
        List(3.14f).map: env =>
          val y = 2.0f
          D(env, f(y))

Notice that all is good but there is no trace of the implementation of f, even though it's marked as inline:

ClassDef(
  "Test",
  DefDef("<init>", List(TermParamClause(Nil)), Inferred(), None),
  List(Apply(Select(New(Inferred()), "<init>"), Nil)),
  None,
  List(
    DefDef(
      "g",
      Nil,
      Inferred(),
      Some(
        Apply(
          TypeApply(
            Select(
              Apply(
                TypeApply(Select(Ident("List"), "apply"), List(Inferred())),
                List(Typed(Repeated(List(Literal(FloatConstant(3.14f))), Inferred()), Inferred()))
              ),
              "map"
            ),
            List(Inferred())
          ),
          List(
            Block(
              List(
                DefDef(
                  "$anonfun",
                  List(TermParamClause(List(ValDef("env", Inferred(), None)))),
                  Inferred(),
                  Some(
                    Block(
                      List(ValDef("y", Inferred(), Some(Literal(FloatConstant(2.0f))))),
                      Apply(
                        Select(Ident("D"), "apply"),
                        List(Ident("env"), Apply(Select(Ident("f"), "apply"), List(Ident("y"))))
                      )
                    )
                  )
                )
              ),
              Closure(Ident("$anonfun"), None)
            )
          )
        )
      )
    )
  )
)

If we move f inside the Test class, then sure enough, the whole definition is available.

That being the case, I don't see how I could write something like this (made up, no idea if it compiles):

    case class C(x: Float, y: Float)
    case class D(x: Float, c: C)

    inline def f: Float => C =
      foo => C(0f, foo)

    @DebugAST2
    class Test:

      def g =
        List(3.14f).map: env =>
          val y = 2.0f
          D(env, f(y))

    @DebugAST2
    class Test2:

      def h =
        List(6f).map: env =>
          D(env, f(1.0f))

Where f is reused between two code blocks and I need to know it's definition in order to transpile and include it in the result of the two macro invocations.

The second best - or maybe even best of all? - would be inline macros... with less aggressive inlining! ...which I think means no val inlining because of the way that function arguments end up getting handled / obliterated. I wonder if there's some workaround there where args are annotated or declared as something that doesn't end up as / being treated like a val? 🤔

[davesmith00000]

Possibly related:

inline val one: 1.0f = 1.0f

inline def foobar(value: Float): vec2 =
  vec2(value, one)
  
DebugAST.anyToAST(foobar(2.0f))

Now looking at that... I'm expecting a tree that sooner or later says "vec2 was applied with (2.0, 1.0)", but alas no... I've got 2.0 and Ident("one"). Any idea why one wasn't inlined?

Inlined(
  None,
  Nil,
  Inlined(
    Some(Apply(Ident("foobar"), List(Literal(FloatConstant(2.0f))))),
    Nil,
    Typed(
      Apply(Select(Ident("vec2"), "apply"), List(Inlined(None, Nil, Literal(FloatConstant(2.0f))), Ident("one"))),
      TypeIdent("vec2")
    )
  )
)

[davesmith00000]

Right, so with thanks to some nice people on Discord...

inline vals are dealt with in a later phase so the macro doesn't get the benefit of them:

...
               inlining  inline and execute macros
           postInlining  add mirror support for inlined code
checkUnusedPostInlining  check for unused elements

... loads of stuff ... then...

             inlineVals  check right hand-sides of an `inline val`s
...

And the desired effect can be achieved with the final keyword:

final val one = 1.0f