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?