Source:
public Expression rewriteLambdaExpression(LambdaExpression node, Context context, ExpressionTreeRewriter<Context> treeRewriter) { LinkedHashSet<Symbol> referencedSymbols = new LinkedHashSet<>(); Expression rewrittenBody = treeRewriter.rewrite(node.getBody(), context.withReferencedSymbols(referencedSymbols)); List<Symbol> lambdaArguments = node.getArguments().stream() .map(LambdaArgumentDeclaration::getName) .map(Identifier::getValue) .map(Symbol::new) .collect(toImmutableList()); referencedSymbols.removeAll(lambdaArguments); Set<Symbol> captureSymbols = referencedSymbols; ImmutableMap.Builder<Symbol, Symbol> captureSymbolToExtraSymbol = ImmutableMap.builder(); ImmutableList.Builder<LambdaArgumentDeclaration> newLambdaArguments = ImmutableList.builder(); for (Symbol captureSymbol : captureSymbols) { Symbol extraSymbol = symbolAllocator.newSymbol(captureSymbol.getName(), symbolTypes.get(captureSymbol)); captureSymbolToExtraSymbol.put(captureSymbol, extraSymbol); newLambdaArguments.add(new LambdaArgumentDeclaration(new Identifier(extraSymbol.getName()))); } newLambdaArguments.addAll(node.getArguments()); ImmutableMap<Symbol, Symbol> symbolsMap = captureSymbolToExtraSymbol.build(); Function<Symbol, Expression> symbolMapping = x -> symbolsMap.getOrDefault(x, x).toSymbolReference(); Expression rewrittenExpression = new LambdaExpression(newLambdaArguments.build(), inlineSymbols(symbolMapping, rewrittenBody)); if (captureSymbols.size() != 0) { List<Expression> capturedValues = captureSymbols.stream() .map(symbol -> new SymbolReference(symbol.getName())) .collect(toImmutableList()); rewrittenExpression = new BindExpression(capturedValues, rewrittenExpression); } context.getReferencedSymbols().addAll(captureSymbols); return rewrittenExpression; }

Target:
public Expression rewriteLambdaExpression(LambdaExpression node, Context context, ExpressionTreeRewriter<Context> treeRewriter) { LinkedHashSet<Symbol> referencedSymbols = new LinkedHashSet<>(); Expression rewrittenBody = treeRewriter.rewrite(node.getBody(), context.withReferencedSymbols(referencedSymbols)); List<Symbol> lambdaArguments = node.getArguments().stream() .map(LambdaArgumentDeclaration::getName) .map(Identifier::getValue) .map(Symbol::new) .collect(toImmutableList()); referencedSymbols.removeAll(lambdaArguments); Set<Symbol> captureSymbols = referencedSymbols; ImmutableMap.Builder<Symbol, Symbol> captureSymbolToExtraSymbol = ImmutableMap.builder(); ImmutableList.Builder<LambdaArgumentDeclaration> newLambdaArguments = ImmutableList.builder(); for (Symbol captureSymbol : captureSymbols) { Symbol extraSymbol = symbolAllocator.newSymbol(captureSymbol.getName(), symbolTypes.get(captureSymbol)); captureSymbolToExtraSymbol.put(captureSymbol, extraSymbol); newLambdaArguments.add(new LambdaArgumentDeclaration(new Identifier(extraSymbol.getName()))); } newLambdaArguments.addAll(node.getArguments()); ImmutableMap<Symbol, Symbol> symbolsMap = captureSymbolToExtraSymbol.build(); Function<Symbol, Expression> symbolMapping = symbol -> symbolsMap.getOrDefault(symbol, symbol).toSymbolReference(); Expression rewrittenExpression = new LambdaExpression(newLambdaArguments.build(), inlineSymbols(symbolMapping, rewrittenBody)); if (captureSymbols.size() != 0) { List<Expression> capturedValues = captureSymbols.stream() .map(symbol -> new SymbolReference(symbol.getName())) .collect(toImmutableList()); rewrittenExpression = new BindExpression(capturedValues, rewrittenExpression); } context.getReferencedSymbols().addAll(captureSymbols); return rewrittenExpression; }