Repository: apple/pkl
File: pkl-core/src/main/java/org/pkl/core/ast/expression/literal/AmendFunctionNode.java (Java)

Code before:
    for (var i = 0; i < hostFrameDesecriptor.getNumberOfSlots(); i++) {
      var slotInfo = hostFrameDesecriptor.getSlotInfo(i);
      // Copy for-generator variables from the outer frame descriptor into inner lambda.
      //
      // We need to do this because at parse time within AstBuilder, we inject for-generator
      // variables into the frame descriptor of the containing root node.
      // The expectation is that when GeneratorForNode executes, it writes for-generator variables
      // into these slots.
      //
      // In the case of an amend function node, AstBuilder can't determine out that there is another
      // frame (e.g. with `new Mixin { ... }` syntax), so it injects for-generator vars into the
      // wrong frame.
      //
      // As a remedy, we simply copy outer for-generator variables into this frame.
      if (slotInfo != null && slotInfo.equals(SymbolTable.FOR_GENERATOR_VARIABLE)) {
        hasForGenVars = true;
        forGenVars.add(
            new Slot(hostFrameDesecriptor.getSlotKind(i), hostFrameDesecriptor.getSlotName(i)));
      } else {
        forGenVars.add(new Slot(hostFrameDesecriptor.getSlotKind(i), Identifier.DUMMY));
      }
    }
    // only add for generator vars if there's any
    if (hasForGenVars) {
      for (var slot : forGenVars) {
        builder.addSlot(slot.kind, slot.name, null);
      }
    }
    var objectToAmendSlot = builder.addSlot(FrameSlotKind.Object, new Object(), null);
    var frameDescriptor = builder.build();

    var subsequentFunctionRootNode =
        new SimpleRootNode(
            hostNode.language,
            frameDescriptor,
            sourceSection,
            hostNode.qualifiedScopeName + ".<function>",
            new AmendFunctionBodyNode(
                sourceSection,
                hostNode.copy(
                    ReadFrameSlotNodeGen.create(
                        hostNode.getParentNode().getSourceSection(), objectToAmendSlot)),
                parameterSlots,
                objectToAmendSlot,
                null));

    if (parameterSlots.length > 0) {
      var parameterCount = hostNode.parameterTypes.length;
      initialFunctionRootNode =
          new FunctionNode(
              hostNode.language,

Review comment (suggestion):
We can avoid the extra allocation here; can just start filling things in the first time we encounter a for generator variable.

```suggestion
        if (!hasForGenVars) {
          hasForGenVars = true;
          for (var j = 0; j < i; j++) {
            builder.addSlot(FrameSlotKind.Illegal, Identifier.DUMMY, null);
          }
        }
        builder.addSlot(hostFrameDesecriptor.getSlotKind(i), hostFrameDesecriptor.getSlotName(i), null);
      } else if (hasForGenVars) {
        builder.addSlot(FrameSlotKind.Illegal, Identifier.DUMMY, null);
```

Code after:
      // Copy for-generator variables from the outer frame descriptor into inner lambda.
      //
      // We need to do this because at parse time within AstBuilder, we inject for-generator
      // variables into the frame descriptor of the containing root node.
      // The expectation is that when GeneratorForNode executes, it writes for-generator variables
      // into these slots.
      //
      // In the case of an amend function node, AstBuilder can't determine out that there is another
      // frame (e.g. with `new Mixin { ... }` syntax), so it injects for-generator vars into the
      // wrong frame.
      //
      // As a remedy, we simply copy outer variables into this frame if there are any for generator
      // variables.
      //
      // We need to preserve the frame slot index, so we insert dummy identifiers
      // for other slots that aren't for generator variables.
      if (slotInfo != null && slotInfo.equals(SymbolTable.FOR_GENERATOR_VARIABLE)) {
        if (!hasForGenVars) {
          hasForGenVars = true;
          for (var j = 0; j < i; j++) {
            builder.addSlot(FrameSlotKind.Illegal, Identifier.DUMMY, null);
          }
        }
        builder.addSlot(
            hostFrameDesecriptor.getSlotKind(i), hostFrameDesecriptor.getSlotName(i), null);
      } else if (hasForGenVars) {
        builder.addSlot(FrameSlotKind.Illegal, Identifier.DUMMY, null);
      }
    }
    var objectToAmendSlot = builder.addSlot(FrameSlotKind.Object, new Object(), null);
    var frameDescriptor = builder.build();

    var subsequentFunctionRootNode =
        new SimpleRootNode(
            hostNode.language,
            frameDescriptor,
            sourceSection,
            hostNode.qualifiedScopeName + ".<function>",
            new AmendFunctionBodyNode(
                sourceSection,
                hostNode.copy(
                    ReadFrameSlotNodeGen.create(
                        hostNode.getParentNode().getSourceSection(), objectToAmendSlot)),
                parameterSlots,
                objectToAmendSlot,
                null));

    if (parameterSlots.length > 0) {
      var parameterCount = hostNode.parameterTypes.length;
      initialFunctionRootNode =
          new FunctionNode(