Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / inline / patch_engine.py: 76%
188 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Apply inline patch IR to copied modeling source."""
17from __future__ import annotations
19import ast
20from collections import defaultdict
21from typing import Iterable
23from hyper_parallel.codegen.astkit.edits import TextEdit, apply_edits
24from hyper_parallel.codegen.astkit.index import build_source_index
25from hyper_parallel.codegen.inline.declaration_edits import declaration_edits
26from hyper_parallel.codegen.inline.forward_edits import forward_extract_edits
27from hyper_parallel.codegen.inline.definition_edits import apply_class_removals
28from hyper_parallel.codegen.inline.ir import (
29 ConstructorReplacePatch,
30 ForwardExtractPatch,
31 InlinePatchSet,
32)
35def apply_patch_set(source_text: str, patch_set: InlinePatchSet) -> str:
36 """Apply imports, constructor edits, class removals, markers, and forward bodies.
38 Args:
39 source_text: Modeling source the patches are applied to.
40 patch_set: Edits emitted by the inline passes.
42 Returns:
43 The patched source text.
44 """
46 text, hits = _apply_source_declarations(source_text, patch_set)
47 text = _apply_constructor_followups(text, patch_set, hits)
48 return _apply_forward_patches(text, patch_set.forward_extracts)
51def _apply_source_declarations(
52 source_text: str, patch_set: InlinePatchSet
53) -> tuple[str, dict[str, int]]:
54 """Apply top-level declarations and the first constructor rewrite wave."""
55 index = build_source_index(source_text)
56 edits = declaration_edits(
57 source_text,
58 index.import_end,
59 index.docstring_end,
60 patch_set.imports,
61 patch_set.module_snippets,
62 )
63 waves = _constructor_waves(list(patch_set.constructor_replaces))
64 hits: dict[str, int] = defaultdict(int)
65 if waves:
66 edits.extend(_constructor_edits(source_text, waves[0], hits))
67 text = apply_edits(source_text, edits) if edits else source_text
68 for wave in waves[1:]:
69 wave_edits = _constructor_edits(text, wave, hits)
70 if wave_edits:
71 text = apply_edits(text, wave_edits)
72 return text, hits
75def _apply_constructor_followups(
76 source_text: str,
77 patch_set: InlinePatchSet,
78 hits: dict[str, int],
79) -> str:
80 """Apply deferred replacements, validate scopes, and remove replaced classes."""
81 deferred_edits = _deferred_replacement_edits(
82 source_text,
83 patch_set.constructor_replaces,
84 hits,
85 )
86 if deferred_edits:
87 source_text = apply_edits(source_text, deferred_edits)
88 if patch_set.constructor_replaces:
89 _require_scoped_rewrites(patch_set.constructor_replaces, hits)
90 if patch_set.class_removals:
91 source_text = apply_class_removals(source_text, patch_set.class_removals)
92 return source_text
95def _apply_forward_patches(
96 source_text: str, patches: Iterable[ForwardExtractPatch]
97) -> str:
98 """Apply forward extraction only after the source structure has settled."""
99 edits = forward_extract_edits(source_text, patches)
100 return apply_edits(source_text, edits) if edits else source_text
103def _constructor_waves(
104 patch_list: list[ConstructorReplacePatch],
105) -> list[list[ConstructorReplacePatch]]:
106 """Group constructor rewrites into dependency-ordered waves.
108 Two compositions force a later wave, re-parsed from the previous
109 wave's output: a patch whose ``old_ctor`` another remaining patch
110 introduces (an ``A -> B`` wrap retargeted further by a ``B -> C``
111 patch), and a ``name`` retarget sharing its ``old_ctor`` with a
112 ``wrap_source`` patch — the inline CP subclass retargets the
113 *outer* wrapper call the wrap introduced, and applying the name
114 swap first would consume the call site the wrap still needs (a
115 scoped wrap with no call site fails loudly, so the order is not
116 optional).
117 """
119 remaining = list(patch_list)
120 waves: list[list[ConstructorReplacePatch]] = []
121 while remaining:
122 introduced = {
123 patch.new_ctor for patch in remaining if patch.new_ctor != patch.old_ctor
124 }
125 wave = [patch for patch in remaining if _is_wave_ready(patch, remaining, introduced)]
126 if not wave:
127 names = sorted({patch.old_ctor for patch in remaining})
128 raise ValueError(
129 f"codegen: constructor replacement dependency cycle among {names}"
130 )
131 chosen = {id(patch) for patch in wave}
132 remaining = [patch for patch in remaining if id(patch) not in chosen]
133 waves.append(wave)
134 return waves
137def _is_wave_ready(
138 patch: ConstructorReplacePatch,
139 remaining: list[ConstructorReplacePatch],
140 introduced: set[str],
141) -> bool:
142 """Return whether a constructor patch has no remaining dependency."""
143 if patch.old_ctor in introduced:
144 return False
145 if patch.mode != "name":
146 return True
147 return not any(
148 other is not patch
149 and other.mode == "wrap_source"
150 and other.old_ctor == patch.old_ctor
151 for other in remaining
152 )
155def _constructor_edits(
156 source_text: str,
157 patches: list[ConstructorReplacePatch],
158 hits: dict[str, int],
159) -> list[TextEdit]:
160 """Build the edits that swap constructor calls per patch target."""
161 patches_by_old: dict[str, list[ConstructorReplacePatch]] = defaultdict(list)
162 for patch in patches:
163 patches_by_old[patch.old_ctor].append(patch)
164 if not patches_by_old:
165 return []
166 scopes = _scopes_by_old_ctor(patches)
167 tree = ast.parse(source_text)
168 offsets = _line_offsets(source_text)
169 edits: list[TextEdit] = []
170 for node in ast.walk(tree):
171 if not isinstance(node, ast.Assign):
172 continue
173 attr = _assigned_attr(node)
174 if attr is None:
175 continue
176 call = node.value
177 if not isinstance(call, ast.Call):
178 continue
179 name = _call_name(call.func)
180 if name is None or name not in patches_by_old:
181 continue
182 scope = scopes[name]
183 if scope is not None and attr not in scope:
184 continue
185 patch = patches_by_old[name][0]
186 hits[name] += 1
187 if patch.mode == "name":
188 start, end = _node_span(call.func, offsets)
189 edits.append(TextEdit(start, end, patch.new_ctor))
190 elif patch.mode == "wrap_source":
191 start, end = _node_span(call, offsets)
192 call_source = source_text[start:end]
193 extra = "".join(f", {arg}" for arg in patch.keyword_args)
194 replacement = f"{patch.new_ctor}(module={call_source}{extra})"
195 edits.append(TextEdit(start, end, replacement))
196 else:
197 raise ValueError(f"unknown constructor replacement mode: {patch.mode!r}")
198 return _drop_nested_edits(edits)
201def _scopes_by_old_ctor(
202 patches: Iterable[ConstructorReplacePatch],
203) -> dict[str, set[str] | None]:
204 """Merge each constructor's rewrite scope; ``None`` means unrestricted."""
206 scopes: dict[str, set[str] | None] = {}
207 for patch in patches:
208 if not patch.scope_attrs:
209 scopes[patch.old_ctor] = None
210 continue
211 merged = scopes.get(patch.old_ctor)
212 if merged is None and patch.old_ctor in scopes:
213 continue
214 if merged is None:
215 merged = scopes[patch.old_ctor] = set()
216 merged.update(patch.scope_attrs)
217 return scopes
220def _assigned_attr(node: ast.Assign) -> str | None:
221 """Return the stored attribute of an attribute store, else ``None``.
223 Module replacements are sunk at the assignment that installs the module, so
224 scoping by the assigned attribute is exactly the FQN granularity the native
225 plan used to decide the replacement. Families attach modules at more than
226 the ``self.<attr>`` store: a loop-local ``layer.<attr>`` upgrade and a
227 ModuleList element ``self.layers[i].<attr>`` install the same matched
228 modules, so any attribute store yields its stored attribute -- the rule's
229 scope set keeps the rewrite at the plan's matched-FQN granularity.
230 """
232 if len(node.targets) != 1:
233 return None
234 target = node.targets[0]
235 return target.attr if isinstance(target, ast.Attribute) else None
238def _require_scoped_rewrites(
239 patches: Iterable[ConstructorReplacePatch],
240 hits: dict[str, int],
241) -> None:
242 """Fail when a scoped rule matched FQNs but no source call site.
244 Silently dropping the rewrite would emit an artifact that no longer matches
245 the plan that decided the replacement, so a missing call site is a
246 generation error: the scope handling needs to learn the new shape rather
247 than the rewrite being widened back to the whole file.
248 """
250 missing = sorted(
251 {patch.old_ctor for patch in patches if patch.scope_attrs and not hits.get(patch.old_ctor)}
252 )
253 if missing:
254 raise RuntimeError(
255 "codegen: no `self.<attr> = <ctor>(...)` call site found for "
256 f"{missing}; the plan replaced modules of this type, so the "
257 "replacement cannot be sunk without diverging from the plan"
258 )
261def _deferred_replacement_edits(
262 source_text: str,
263 patches: Iterable[ConstructorReplacePatch],
264 hits: dict[str, int],
265) -> list[TextEdit]:
266 """Install observed wrappers after an imported owner built its subtree.
268 Some source models inherit a decoder whose constructor lives in an
269 installed dependency. The copied modeling unit therefore has no direct
270 constructor call to rewrite. For a wrap-style replacement, the native
271 operation is still expressible exactly: wait until the root model has
272 built the observed FQNs, then replace those module attributes in place.
273 """
274 pending = [
275 patch
276 for patch in patches
277 if patch.scope_fqns and not hits.get(patch.old_ctor)
278 ]
279 if not pending:
280 return []
281 tree = ast.parse(source_text)
282 offsets = _line_offsets(source_text)
283 insertions: dict[int, list[str]] = defaultdict(list)
284 for patch in pending:
285 if patch.mode != "wrap_source":
286 continue
287 root_attr = patch.scope_fqns[0].split(".", 1)[0]
288 constructors = _root_constructors(tree, root_attr)
289 if len(constructors) != 1:
290 continue
291 function = constructors[0]
292 _, end = _node_span(function.body[-1], offsets)
293 insertions[end].append(_render_deferred_replacement(patch))
294 hits[patch.old_ctor] += 1
295 return [
296 TextEdit(offset, offset, "".join(parts))
297 for offset, parts in insertions.items()
298 ]
301def _root_constructors(tree: ast.Module, root_attr: str) -> list[ast.FunctionDef]:
302 """Return constructors that establish the first segment of an observed FQN."""
303 constructors = []
304 for node in tree.body:
305 if not isinstance(node, ast.ClassDef):
306 continue
307 for child in node.body:
308 if not isinstance(child, ast.FunctionDef) or child.name != "__init__":
309 continue
310 if any(
311 isinstance(target, ast.Attribute)
312 and isinstance(target.value, ast.Name)
313 and target.value.id == "self"
314 and target.attr == root_attr
315 for item in ast.walk(child)
316 if isinstance(item, (ast.Assign, ast.AnnAssign))
317 for target in (
318 item.targets if isinstance(item, ast.Assign) else (item.target,)
319 )
320 ):
321 constructors.append(child)
322 return constructors
325def _render_deferred_replacement(patch: ConstructorReplacePatch) -> str:
326 """Render an exact-FQN replacement loop inside a root constructor."""
327 arguments = [
328 argument
329 for argument in patch.keyword_args
330 if not argument.startswith("module_fqn=")
331 ]
332 extra = "".join(f", {argument}" for argument in arguments)
333 fqns = repr(tuple(patch.scope_fqns))
334 return (
335 "\n _hp_replacement_cache = {}\n"
336 f" for _hp_fqn in {fqns}:\n"
337 " _hp_parent_fqn, _, _hp_name = _hp_fqn.rpartition('.')\n"
338 " _hp_parent = self.get_submodule(_hp_parent_fqn)\n"
339 " _hp_source = getattr(_hp_parent, _hp_name)\n"
340 " _hp_key = id(_hp_source)\n"
341 " if _hp_key not in _hp_replacement_cache:\n"
342 f" _hp_replacement_cache[_hp_key] = {patch.new_ctor}(\n"
343 f" module=_hp_source, module_fqn=_hp_fqn{extra})\n"
344 " setattr(_hp_parent, _hp_name, _hp_replacement_cache[_hp_key])"
345 )
348def _drop_nested_edits(edits: list[TextEdit]) -> list[TextEdit]:
349 """Keep outer constructor rewrites when constructor calls are nested."""
350 kept: list[TextEdit] = []
351 covered_end = -1
352 for edit in sorted(edits, key=lambda item: (item.start, -item.end)):
353 if edit.start < covered_end:
354 continue
355 kept.append(edit)
356 covered_end = edit.end
357 return kept
360def _call_name(node: ast.AST) -> str | None:
361 if isinstance(node, ast.Name):
362 return node.id
363 if isinstance(node, ast.Attribute):
364 return node.attr
365 return None
368def _node_span(node: ast.AST, offsets: tuple[int, ...]) -> tuple[int, int]:
369 """Return the ``(start, end)`` source span of ``node``.
371 Decorators are part of the span: deleting only the ``def``/``class`` line
372 would leave a dangling decorator that then applies to whatever statement
373 follows the removed block.
374 """
376 if (
377 getattr(node, "lineno", None) is None
378 or getattr(node, "col_offset", None) is None
379 or getattr(node, "end_lineno", None) is None
380 or getattr(node, "end_col_offset", None) is None
381 ):
382 raise ValueError(f"AST node has no source span: {type(node).__name__}")
383 line = node.lineno
384 column = node.col_offset
385 decorators = getattr(node, "decorator_list", None)
386 if decorators:
387 first = min(decorators, key=lambda item: item.lineno)
388 if first.lineno < line:
389 # A decorator node starts at its expression, one column *after* the
390 # ``@``; the span must start at the ``@`` itself, otherwise deleting
391 # the block leaves a stray ``@`` behind.
392 line, column = first.lineno, max(0, first.col_offset - 1)
393 start = offsets[line - 1] + column
394 end = offsets[node.end_lineno - 1] + node.end_col_offset
395 return start, end
398def _line_offsets(source_text: str) -> tuple[int, ...]:
399 offsets = [0]
400 total = 0
401 for line in source_text.splitlines(keepends=True):
402 total += len(line)
403 offsets.append(total)
404 return tuple(offsets)