@openrewrite/rewrite 8.66.1 → 8.66.2

package/dist/javascript/templating/engine.js

@@ -10,6 +10,8 @@ var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, ge
 };
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.TemplateApplier = exports.TemplateEngine = void 0;
+exports.setTemplateSourceFileCache = setTemplateSourceFileCache;
+exports.clearTemplateCache = clearTemplateCache;
 /*
  * Copyright 2025 the original author or authors.
  * <p>
@@ -32,73 +34,185 @@ const immer_1 = require("immer");
 const utils_1 = require("./utils");
 const capture_1 = require("./capture");
 const placeholder_replacement_1 = require("./placeholder-replacement");
+const format_1 = require("../format");
+const parser_utils_1 = require("../parser-utils");
+const uuid_1 = require("../../uuid");
+const dependency_workspace_1 = require("../dependency-workspace");
 /**
- * Cache for compiled templates.
+ * Simple LRU (Least Recently Used) cache implementation.
+ * Used for template/pattern compilation caching with bounded memory usage.
  */
-const templateCache = new utils_1.TemplateCache();
+class LRUCache {
+    constructor(maxSize) {
+        this.maxSize = maxSize;
+        this.cache = new Map();
+    }
+    get(key) {
+        const value = this.cache.get(key);
+        if (value !== undefined) {
+            // Move to end (most recently used)
+            this.cache.delete(key);
+            this.cache.set(key, value);
+        }
+        return value;
+    }
+    set(key, value) {
+        // Remove if exists (to update position)
+        this.cache.delete(key);
+        // Add to end
+        this.cache.set(key, value);
+        // Evict oldest if over capacity
+        if (this.cache.size > this.maxSize) {
+            const iterator = this.cache.keys();
+            const firstEntry = iterator.next();
+            if (!firstEntry.done) {
+                this.cache.delete(firstEntry.value);
+            }
+        }
+    }
+    clear() {
+        this.cache.clear();
+    }
+}
+/**
+ * Module-level TypeScript sourceFileCache for template parsing.
+ */
+let templateSourceFileCache;
+/**
+ * Configure the sourceFileCache used for template parsing.
+ *
+ * @param cache The sourceFileCache to use, or undefined to disable caching
+ */
+function setTemplateSourceFileCache(cache) {
+    templateSourceFileCache = cache;
+}
+/**
+ * Cache for compiled templates and patterns.
+ * Stores parsed ASTs to avoid expensive re-parsing and dependency resolution.
+ * Bounded to 100 entries using LRU eviction to prevent unbounded memory growth.
+ */
+class TemplateCache {
+    constructor() {
+        this.cache = new LRUCache(100);
+    }
+    /**
+     * Generates a cache key from template string, captures, and options.
+     */
+    generateKey(templateString, captures, contextStatements, dependencies) {
+        // Use the actual template string (with placeholders) as the primary key
+        const templateKey = templateString;
+        // Capture names
+        const capturesKey = captures.map(c => c.getName()).join(',');
+        // Context statements
+        const contextKey = contextStatements.join(';');
+        // Dependencies
+        const depsKey = JSON.stringify(dependencies || {});
+        return `${templateKey}::${capturesKey}::${contextKey}::${depsKey}`;
+    }
+    /**
+     * Gets a cached compilation unit or creates and caches a new one.
+     */
+    getOrParse(templateString, captures, contextStatements, dependencies) {
+        return __awaiter(this, void 0, void 0, function* () {
+            const key = this.generateKey(templateString, captures, contextStatements, dependencies);
+            let cu = this.cache.get(key);
+            if (cu) {
+                return cu;
+            }
+            // Create workspace if dependencies are provided
+            // DependencyWorkspace has its own cache, so multiple templates with
+            // the same dependencies will automatically share the same workspace
+            let workspaceDir;
+            if (dependencies && Object.keys(dependencies).length > 0) {
+                workspaceDir = yield dependency_workspace_1.DependencyWorkspace.getOrCreateWorkspace(dependencies);
+            }
+            // Prepend context statements for type attribution context
+            const fullTemplateString = contextStatements.length > 0
+                ? contextStatements.join('\n') + '\n' + templateString
+                : templateString;
+            // Parse and cache (workspace only needed during parsing)
+            // Use templateSourceFileCache if configured for ~3.2x speedup on dependency file parsing
+            const parser = new __2.JavaScriptParser({
+                relativeTo: workspaceDir,
+                sourceFileCache: templateSourceFileCache
+            });
+            const parseGenerator = parser.parse({ text: fullTemplateString, sourcePath: 'template.ts' });
+            cu = (yield parseGenerator.next()).value;
+            this.cache.set(key, cu);
+            return cu;
+        });
+    }
+    /**
+     * Clears the cache.
+     */
+    clear() {
+        this.cache.clear();
+    }
+}
+/**
+ * Cache for compiled templates and patterns.
+ * Private to the engine module - encapsulates caching implementation.
+ */
+const templateCache = new TemplateCache();
+/**
+ * Clears the template cache. Only exported for testing and benchmarking purposes.
+ * Normal application code should not need to call this.
+ */
+function clearTemplateCache() {
+    templateCache.clear();
+}
 /**
  * Internal template engine - handles the core templating logic.
  * Not exported from index, so only visible within the templating module.
  */
 class TemplateEngine {
     /**
-     * Applies a template with optional match results from pattern matching.
+     * Gets the parsed and extracted template tree (before value substitution).
+     * This is the cacheable part of template processing.
      *
      * @param templateParts The string parts of the template
      * @param parameters The parameters between the string parts
-     * @param cursor The cursor pointing to the current location in the AST
-     * @param coordinates The coordinates specifying where and how to insert the generated AST
-     * @param values Map of capture names to values to replace the parameters with
-     * @param wrappersMap Map of capture names to J.RightPadded wrappers (for preserving markers)
      * @param contextStatements Context declarations (imports, types, etc.) to prepend for type attribution
      * @param dependencies NPM dependencies for type attribution
-     * @returns A Promise resolving to the generated AST node
+     * @returns A Promise resolving to the extracted template AST
      */
-    static applyTemplate(templateParts_1, parameters_1, cursor_1, coordinates_1) {
-        return __awaiter(this, arguments, void 0, function* (templateParts, parameters, cursor, coordinates, values = new Map(), wrappersMap = new Map(), contextStatements = [], dependencies = {}) {
+    static getTemplateTree(templateParts_1, parameters_1) {
+        return __awaiter(this, arguments, void 0, function* (templateParts, parameters, contextStatements = [], dependencies = {}) {
+            // Generate type preamble for captures/parameters with types
+            const preamble = TemplateEngine.generateTypePreamble(parameters);
             // Build the template string with parameter placeholders
             const templateString = TemplateEngine.buildTemplateString(templateParts, parameters);
-            // If the template string is empty, return undefined
-            if (!templateString.trim()) {
-                return undefined;
-            }
+            // Add preamble to context statements (so they're skipped during extraction)
+            const contextWithPreamble = preamble.length > 0
+                ? [...contextStatements, ...preamble]
+                : contextStatements;
             // Use cache to get or parse the compilation unit
-            // For templates, we don't have captures, so use empty array
-            const cu = yield templateCache.getOrParse(templateString, [], // templates don't have captures in the cache key
-            contextStatements, dependencies);
+            const cu = yield templateCache.getOrParse(templateString, [], contextWithPreamble, dependencies);
             // Check if there are any statements
             if (!cu.statements || cu.statements.length === 0) {
                 throw new Error(`Failed to parse template code (no statements):\n${templateString}`);
             }
-            // Skip context statements to get to the actual template code
-            const templateStatementIndex = contextStatements.length;
-            if (templateStatementIndex >= cu.statements.length) {
-                return undefined;
-            }
-            // Extract the relevant part of the AST
-            const firstStatement = cu.statements[templateStatementIndex].element;
-            let extracted;
-            // Check if this is a wrapped template (function __TEMPLATE__() { ... })
-            if (firstStatement.kind === java_1.J.Kind.MethodDeclaration) {
-                const func = firstStatement;
-                if (func.name.simpleName === '__TEMPLATE__' && func.body) {
-                    // __TEMPLATE__ wrapper indicates the original template was a block.
-                    // Always return the block to preserve the block structure.
-                    extracted = func.body;
-                }
-                else {
-                    // Not a __TEMPLATE__ wrapper
-                    extracted = firstStatement;
-                }
-            }
-            else if (firstStatement.kind === __2.JS.Kind.ExpressionStatement) {
-                extracted = firstStatement.expression;
-            }
-            else {
-                extracted = firstStatement;
-            }
-            // Create a copy to avoid sharing cached AST instances
-            const ast = (0, immer_1.produce)(extracted, _ => { });
+            // The template code is always the last statement (after context + preamble)
+            const lastStatement = cu.statements[cu.statements.length - 1].element;
+            // Extract from wrapper using shared utility
+            const extracted = utils_1.PlaceholderUtils.extractFromWrapper(lastStatement, 'Template');
+            return (0, immer_1.produce)(extracted, _ => { });
+        });
+    }
+    /**
+     * Applies a template from a pre-parsed AST and returns the resulting AST.
+     * This method is used by Template.apply() after getting the cached template tree.
+     *
+     * @param ast The pre-parsed template AST
+     * @param parameters The parameters between the string parts
+     * @param cursor The cursor pointing to the current location in the AST
+     * @param coordinates The coordinates specifying where and how to insert the generated AST
+     * @param values Map of capture names to values to replace the parameters with
+     * @param wrappersMap Map of capture names to J.RightPadded wrappers (for preserving markers)
+     * @returns A Promise resolving to the generated AST node
+     */
+    static applyTemplateFromAst(ast_1, parameters_1, cursor_1, coordinates_1) {
+        return __awaiter(this, arguments, void 0, function* (ast, parameters, cursor, coordinates, values = new Map(), wrappersMap = new Map()) {
             // Create substitutions map for placeholders
             const substitutions = new Map();
             for (let i = 0; i < parameters.length; i++) {
@@ -112,8 +226,68 @@ class TemplateEngine {
             return new TemplateApplier(cursor, coordinates, unsubstitutedAst).apply();
         });
     }
+    /**
+     * Generates type preamble declarations for captures/parameters with type annotations.
+     *
+     * @param parameters The parameters
+     * @returns Array of preamble statements
+     */
+    static generateTypePreamble(parameters) {
+        const preamble = [];
+        for (let i = 0; i < parameters.length; i++) {
+            const param = parameters[i].value;
+            const placeholder = `${utils_1.PlaceholderUtils.PLACEHOLDER_PREFIX}${i}__`;
+            // Check for Capture (could be a Proxy, so check for symbol property)
+            const isCapture = param instanceof capture_1.CaptureImpl ||
+                (param && typeof param === 'object' && param[capture_1.CAPTURE_NAME_SYMBOL]);
+            const isCaptureValue = param instanceof capture_1.CaptureValue;
+            const isTreeArray = Array.isArray(param) && param.length > 0 && (0, __1.isTree)(param[0]);
+            if (isCapture) {
+                const captureType = param[capture_1.CAPTURE_TYPE_SYMBOL];
+                if (captureType) {
+                    const typeString = typeof captureType === 'string'
+                        ? captureType
+                        : this.typeToString(captureType);
+                    // Only add preamble if we have a concrete type (not 'any')
+                    if (typeString !== 'any') {
+                        preamble.push(`let ${placeholder}: ${typeString};`);
+                    }
+                }
+            }
+            else if (isCaptureValue) {
+                // For CaptureValue, check if the root capture has a type
+                const rootCapture = param.rootCapture;
+                if (rootCapture) {
+                    const captureType = rootCapture[capture_1.CAPTURE_TYPE_SYMBOL];
+                    if (captureType) {
+                        const typeString = typeof captureType === 'string'
+                            ? captureType
+                            : this.typeToString(captureType);
+                        // Only add preamble if we have a concrete type (not 'any')
+                        if (typeString !== 'any') {
+                            preamble.push(`let ${placeholder}: ${typeString};`);
+                        }
+                    }
+                }
+            }
+            else if ((0, __1.isTree)(param) && !isTreeArray) {
+                // For J elements, derive type from the element's type property if it exists
+                const jElement = param;
+                if (jElement.type) {
+                    const typeString = this.typeToString(jElement.type);
+                    // Only add preamble if we have a concrete type (not 'any')
+                    if (typeString !== 'any') {
+                        preamble.push(`let ${placeholder}: ${typeString};`);
+                    }
+                }
+            }
+        }
+        return preamble;
+    }
     /**
      * Builds a template string with parameter placeholders.
+     * RawCode parameters are spliced directly into the template at construction time.
+     * Other parameters use placeholders that are replaced during application.
      *
      * @param templateParts The string parts of the template
      * @param parameters The parameters between the string parts
@@ -125,32 +299,258 @@ class TemplateEngine {
             result += templateParts[i];
             if (i < parameters.length) {
                 const param = parameters[i].value;
-                // Use a placeholder for Captures, TemplateParams, CaptureValues, Tree nodes, and Tree arrays
-                // Inline everything else (strings, numbers, booleans) directly
-                // Check for Capture (could be a Proxy, so check for symbol property)
-                const isCapture = param instanceof capture_1.CaptureImpl ||
-                    (param && typeof param === 'object' && param[capture_1.CAPTURE_NAME_SYMBOL]);
-                const isTemplateParam = param instanceof capture_1.TemplateParamImpl;
-                const isCaptureValue = param instanceof capture_1.CaptureValue;
-                const isTreeArray = Array.isArray(param) && param.length > 0 && (0, __1.isTree)(param[0]);
-                if (isCapture || isTemplateParam || isCaptureValue || (0, __1.isTree)(param) || isTreeArray) {
-                    const placeholder = `${utils_1.PlaceholderUtils.PLACEHOLDER_PREFIX}${i}__`;
-                    result += placeholder;
+                // Check if this is a RawCode instance - splice directly
+                if (param instanceof capture_1.RawCode || (param && typeof param === 'object' && param[capture_1.RAW_CODE_SYMBOL])) {
+                    result += param.code;
                 }
                 else {
-                    result += param;
+                    // All other parameters use placeholders
+                    // This ensures templates with the same structure always produce the same AST
+                    const placeholder = `${utils_1.PlaceholderUtils.PLACEHOLDER_PREFIX}${i}__`;
+                    result += placeholder;
                 }
             }
         }
-        // Detect if this is a block template that needs wrapping
-        const trimmed = result.trim();
-        if (trimmed.startsWith('{') && trimmed.endsWith('}')) {
-            result = `function __TEMPLATE__() ${result}`;
+        // Always wrap in function body - let the parser decide what it is,
+        // then we'll extract intelligently based on what was parsed
+        return `function ${utils_1.WRAPPER_FUNCTION_NAME}() { ${result} }`;
+    }
+    /**
+     * Converts a Type instance to a TypeScript type string.
+     *
+     * @param type The Type instance
+     * @returns A TypeScript type string
+     */
+    static typeToString(type) {
+        // Handle Type.Class and Type.ShallowClass - return their fully qualified names
+        if (type.kind === java_1.Type.Kind.Class || type.kind === java_1.Type.Kind.ShallowClass) {
+            const classType = type;
+            return classType.fullyQualifiedName;
         }
-        return result;
+        // Handle Type.Primitive - map to TypeScript primitive types
+        if (type.kind === java_1.Type.Kind.Primitive) {
+            const primitiveType = type;
+            switch (primitiveType.keyword) {
+                case 'String':
+                    return 'string';
+                case 'boolean':
+                    return 'boolean';
+                case 'double':
+                case 'float':
+                case 'int':
+                case 'long':
+                case 'short':
+                case 'byte':
+                    return 'number';
+                case 'void':
+                    return 'void';
+                default:
+                    return 'any';
+            }
+        }
+        // Handle Type.Array - render component type plus []
+        if (type.kind === java_1.Type.Kind.Array) {
+            const arrayType = type;
+            const componentTypeString = this.typeToString(arrayType.elemType);
+            return `${componentTypeString}[]`;
+        }
+        // For other types, return 'any' as a fallback
+        // TODO: Implement proper Type to string conversion for other Type.Kind values
+        return 'any';
+    }
+    /**
+     * Gets the parsed and extracted pattern tree with capture markers attached.
+     * This is the entry point for pattern processing, providing pattern-specific
+     * functionality on top of the shared template tree generation.
+     *
+     * @param templateParts The string parts of the template
+     * @param captures The captures between the string parts (can include RawCode)
+     * @param contextStatements Context declarations (imports, types, etc.) to prepend for type attribution
+     * @param dependencies NPM dependencies for type attribution
+     * @returns A Promise resolving to the extracted pattern AST with capture markers
+     */
+    static getPatternTree(templateParts_1, captures_1) {
+        return __awaiter(this, arguments, void 0, function* (templateParts, captures, contextStatements = [], dependencies = {}) {
+            // Generate type preamble for captures with types (skip RawCode)
+            const preamble = [];
+            for (const capture of captures) {
+                // Skip raw code - it's not a capture
+                if (capture instanceof capture_1.RawCode || (capture && typeof capture === 'object' && capture[capture_1.RAW_CODE_SYMBOL])) {
+                    continue;
+                }
+                const captureName = capture[capture_1.CAPTURE_NAME_SYMBOL] || capture.getName();
+                const captureType = capture[capture_1.CAPTURE_TYPE_SYMBOL];
+                if (captureType) {
+                    // Convert Type to string if needed
+                    const typeString = typeof captureType === 'string'
+                        ? captureType
+                        : this.typeToString(captureType);
+                    // Only add preamble if we have a concrete type (not 'any')
+                    if (typeString !== 'any') {
+                        const placeholder = utils_1.PlaceholderUtils.createCapture(captureName, undefined);
+                        preamble.push(`let ${placeholder}: ${typeString};`);
+                    }
+                }
+                // Don't add preamble declarations without types - they don't provide type attribution
+            }
+            // Build the template string with placeholders for captures and raw code
+            let result = '';
+            for (let i = 0; i < templateParts.length; i++) {
+                result += templateParts[i];
+                if (i < captures.length) {
+                    const capture = captures[i];
+                    // Check if this is a RawCode instance - splice directly
+                    if (capture instanceof capture_1.RawCode || (capture && typeof capture === 'object' && capture[capture_1.RAW_CODE_SYMBOL])) {
+                        result += capture.code;
+                    }
+                    else {
+                        // Use symbol to access capture name without triggering Proxy
+                        const captureName = capture[capture_1.CAPTURE_NAME_SYMBOL] || capture.getName();
+                        result += utils_1.PlaceholderUtils.createCapture(captureName, undefined);
+                    }
+                }
+            }
+            // Always wrap in function body - let the parser decide what it is,
+            // then we'll extract intelligently based on what was parsed
+            const templateString = `function ${utils_1.WRAPPER_FUNCTION_NAME}() { ${result} }`;
+            // Add preamble to context statements (so they're skipped during extraction)
+            const contextWithPreamble = preamble.length > 0
+                ? [...contextStatements, ...preamble]
+                : contextStatements;
+            // Filter out RawCode from captures for cache and marker attachment
+            const actualCaptures = captures.filter(c => !(c instanceof capture_1.RawCode || (c && typeof c === 'object' && c[capture_1.RAW_CODE_SYMBOL])));
+            // Use cache to get or parse the compilation unit
+            const cu = yield templateCache.getOrParse(templateString, actualCaptures, contextWithPreamble, dependencies);
+            // Check if there are any statements
+            if (!cu.statements || cu.statements.length === 0) {
+                throw new Error(`Failed to parse pattern code (no statements):\n${templateString}`);
+            }
+            // The pattern code is always the last statement (after context + preamble)
+            const lastStatement = cu.statements[cu.statements.length - 1].element;
+            // Extract from wrapper using shared utility
+            const extracted = utils_1.PlaceholderUtils.extractFromWrapper(lastStatement, 'Pattern');
+            // Attach CaptureMarkers to capture identifiers (only for actual captures, not raw code)
+            const visitor = new MarkerAttachmentVisitor(actualCaptures);
+            return (yield visitor.visit(extracted, undefined));
+        });
     }
 }
 exports.TemplateEngine = TemplateEngine;
+/**
+ * Visitor that attaches CaptureMarkers to capture identifiers in pattern ASTs.
+ * This allows efficient capture detection without string parsing during matching.
+ * Used by TemplateEngine.getPatternTree() for pattern-specific processing.
+ */
+class MarkerAttachmentVisitor extends __2.JavaScriptVisitor {
+    constructor(captures) {
+        super();
+        this.captures = captures;
+    }
+    /**
+     * Attaches CaptureMarker to capture identifiers.
+     */
+    visitIdentifier(ident, p) {
+        const _super = Object.create(null, {
+            visitIdentifier: { get: () => super.visitIdentifier }
+        });
+        return __awaiter(this, void 0, void 0, function* () {
+            var _a, _b;
+            // First call parent to handle standard visitation
+            const visited = yield _super.visitIdentifier.call(this, ident, p);
+            if (!visited || visited.kind !== java_1.J.Kind.Identifier) {
+                return visited;
+            }
+            ident = visited;
+            // Check if this is a capture placeholder
+            if ((_a = ident.simpleName) === null || _a === void 0 ? void 0 : _a.startsWith(utils_1.PlaceholderUtils.CAPTURE_PREFIX)) {
+                const captureInfo = utils_1.PlaceholderUtils.parseCapture(ident.simpleName);
+                if (captureInfo) {
+                    // Find the original capture object to get variadic options and constraint
+                    const captureObj = this.captures.find(c => c.getName() === captureInfo.name);
+                    const variadicOptions = captureObj === null || captureObj === void 0 ? void 0 : captureObj.getVariadicOptions();
+                    const constraint = (_b = captureObj === null || captureObj === void 0 ? void 0 : captureObj.getConstraint) === null || _b === void 0 ? void 0 : _b.call(captureObj);
+                    // Add CaptureMarker to the Identifier with constraint
+                    const marker = new utils_1.CaptureMarker(captureInfo.name, variadicOptions, constraint);
+                    return (0, __1.updateIfChanged)(ident, {
+                        markers: Object.assign(Object.assign({}, ident.markers), { markers: [...ident.markers.markers, marker] })
+                    });
+                }
+            }
+            return ident;
+        });
+    }
+    /**
+     * Propagates markers from element to RightPadded wrapper.
+     */
+    visitRightPadded(right, p) {
+        return __awaiter(this, void 0, void 0, function* () {
+            if (!(0, __1.isTree)(right.element)) {
+                return right;
+            }
+            const visitedElement = yield this.visit(right.element, p);
+            if (visitedElement && visitedElement !== right.element) {
+                const result = yield (0, __1.produceAsync)(right, (draft) => __awaiter(this, void 0, void 0, function* () {
+                    // Visit element first
+                    if (right.element && right.element.kind) {
+                        // Check if element has a CaptureMarker
+                        const elementMarker = utils_1.PlaceholderUtils.getCaptureMarker(visitedElement);
+                        if (elementMarker) {
+                            draft.markers.markers.push(elementMarker);
+                        }
+                        else {
+                            draft.element = visitedElement;
+                        }
+                    }
+                }));
+                return result;
+            }
+            return right;
+        });
+    }
+    /**
+     * Propagates markers from expression to ExpressionStatement.
+     */
+    visitExpressionStatement(expressionStatement, p) {
+        return __awaiter(this, void 0, void 0, function* () {
+            // Visit the expression
+            const visitedExpression = yield this.visit(expressionStatement.expression, p);
+            // Check if expression has a CaptureMarker
+            const expressionMarker = utils_1.PlaceholderUtils.getCaptureMarker(visitedExpression);
+            if (expressionMarker) {
+                return (0, __1.updateIfChanged)(expressionStatement, {
+                    markers: Object.assign(Object.assign({}, expressionStatement.markers), { markers: [...expressionStatement.markers.markers, expressionMarker] }),
+                });
+            }
+            // No marker to move, just update with visited expression
+            return (0, __1.updateIfChanged)(expressionStatement, {
+                expression: visitedExpression
+            });
+        });
+    }
+    /**
+     * Propagates markers from name identifier to BindingElement.
+     * This handles destructuring patterns like {${props}} where the capture marker
+     * is on the identifier but needs to be on the BindingElement for container matching.
+     */
+    visitBindingElement(bindingElement, p) {
+        return __awaiter(this, void 0, void 0, function* () {
+            // Visit the name
+            const visitedName = yield this.visit(bindingElement.name, p);
+            // Check if name has a CaptureMarker
+            const nameMarker = utils_1.PlaceholderUtils.getCaptureMarker(visitedName);
+            if (nameMarker) {
+                return (0, __1.updateIfChanged)(bindingElement, {
+                    name: visitedName,
+                    markers: Object.assign(Object.assign({}, bindingElement.markers), { markers: [...bindingElement.markers.markers, nameMarker] }),
+                });
+            }
+            // No marker to move, just update with visited name
+            return (0, __1.updateIfChanged)(bindingElement, {
+                name: visitedName
+            });
+        });
+    }
+}
 /**
  * Helper class for applying a template to an AST.
  */
@@ -171,11 +571,9 @@ class TemplateApplier {
             // Apply the template based on the location and mode
             switch (loc || 'EXPRESSION_PREFIX') {
                 case 'EXPRESSION_PREFIX':
-                    return this.applyToExpression();
                 case 'STATEMENT_PREFIX':
-                    return this.applyToStatement();
                 case 'BLOCK_END':
-                    return this.applyToBlock();
+                    return this.applyInternal();
                 default:
                     throw new Error(`Unsupported location: ${loc}`);
             }
@@ -186,42 +584,77 @@ class TemplateApplier {
      *
      * @returns A Promise resolving to the modified AST
      */
-    applyToExpression() {
+    applyInternal() {
         return __awaiter(this, void 0, void 0, function* () {
             const { tree } = this.coordinates;
-            // Create a copy of the AST with the prefix from the target
-            return tree ? (0, immer_1.produce)(this.ast, draft => {
-                draft.prefix = tree.prefix;
-            }) : this.ast;
+            if (!tree) {
+                return this.ast;
+            }
+            const originalTree = tree;
+            const resultToUse = this.wrapTree(originalTree, this.ast);
+            return this.format(resultToUse, originalTree);
         });
     }
-    /**
-     * Applies the template to a statement.
-     *
-     * @returns A Promise resolving to the modified AST
-     */
-    applyToStatement() {
+    format(resultToUse, originalTree) {
         return __awaiter(this, void 0, void 0, function* () {
-            const { tree } = this.coordinates;
+            var _a;
             // Create a copy of the AST with the prefix from the target
-            return (0, immer_1.produce)(this.ast, draft => {
-                draft.prefix = tree.prefix;
-            });
+            const result = Object.assign(Object.assign({}, resultToUse), { 
+                // We temporarily set the ID so that the formatter can identify the tree
+                id: originalTree.id, prefix: originalTree.prefix });
+            // Apply auto-formatting to the result
+            const formatted = yield (0, format_1.maybeAutoFormat)(originalTree, result, null, undefined, (_a = this.cursor) === null || _a === void 0 ? void 0 : _a.parent);
+            // Restore the original ID
+            return Object.assign(Object.assign({}, formatted), { id: resultToUse.id });
         });
     }
-    /**
-     * Applies the template to a block.
-     *
-     * @returns A Promise resolving to the modified AST
-     */
-    applyToBlock() {
-        return __awaiter(this, void 0, void 0, function* () {
-            const { tree } = this.coordinates;
-            // Create a copy of the AST with the prefix from the target
-            return (0, immer_1.produce)(this.ast, draft => {
-                draft.prefix = tree.prefix;
-            });
-        });
+    wrapTree(originalTree, resultToUse) {
+        var _a, _b;
+        const parentTree = (_b = (_a = this.cursor) === null || _a === void 0 ? void 0 : _a.parentTree()) === null || _b === void 0 ? void 0 : _b.value;
+        // Only apply wrapping logic if we have parent context
+        if (parentTree) {
+            // FIXME: This is a heuristic to determine if the parent expects a statement child
+            const parentExpectsStatement = parentTree.kind === java_1.J.Kind.Block ||
+                parentTree.kind === java_1.J.Kind.Case ||
+                parentTree.kind === java_1.J.Kind.DoWhileLoop ||
+                parentTree.kind === java_1.J.Kind.ForEachLoop ||
+                parentTree.kind === java_1.J.Kind.ForLoop ||
+                parentTree.kind === java_1.J.Kind.If ||
+                parentTree.kind === java_1.J.Kind.IfElse ||
+                parentTree.kind === java_1.J.Kind.WhileLoop ||
+                parentTree.kind === __2.JS.Kind.CompilationUnit ||
+                parentTree.kind === __2.JS.Kind.ForInLoop;
+            const originalIsStatement = (0, parser_utils_1.isStatement)(originalTree);
+            const resultIsStatement = (0, parser_utils_1.isStatement)(resultToUse);
+            const resultIsExpression = (0, parser_utils_1.isExpression)(resultToUse);
+            // Determine context and wrap if needed
+            if (parentExpectsStatement && originalIsStatement) {
+                // Statement context: wrap in ExpressionStatement if result is not a statement
+                if (!resultIsStatement && resultIsExpression) {
+                    resultToUse = {
+                        kind: __2.JS.Kind.ExpressionStatement,
+                        id: (0, uuid_1.randomId)(),
+                        prefix: resultToUse.prefix,
+                        markers: resultToUse.markers,
+                        expression: Object.assign(Object.assign({}, resultToUse), { prefix: java_1.emptySpace })
+                    };
+                }
+            }
+            else if (!parentExpectsStatement) {
+                // Expression context: wrap in StatementExpression if result is statement-only
+                if (resultIsStatement && !resultIsExpression) {
+                    const stmt = resultToUse;
+                    resultToUse = {
+                        kind: __2.JS.Kind.StatementExpression,
+                        id: (0, uuid_1.randomId)(),
+                        prefix: stmt.prefix,
+                        markers: stmt.markers,
+                        statement: Object.assign(Object.assign({}, stmt), { prefix: java_1.emptySpace })
+                    };
+                }
+            }
+        }
+        return resultToUse;
     }
 }
 exports.TemplateApplier = TemplateApplier;