Module refinery.lib.scripts.ps1.deobfuscation.constants
Inline constant variable references in PowerShell scripts.
Expand source code Browse git
"""
Inline constant variable references in PowerShell scripts.
"""
from __future__ import annotations
from collections import defaultdict
from typing import Iterator
from refinery.lib.scripts import (
Expression,
Node,
Transformer,
_clone_node,
)
from refinery.lib.scripts.ps1.analysis.cache import model_cache
from refinery.lib.scripts.ps1.analysis.dataflow import Ps1VariableFlow
from refinery.lib.scripts.ps1.analysis.model import (
Binding,
binding_key,
is_assignment_write_target,
is_substitutable_position,
is_write_occurrence,
Ps1OccurrenceRole,
)
from refinery.lib.scripts.ps1.analysis.values import unwrap_to_array_literal
from refinery.lib.scripts.ps1.ast import (
assignment_of,
assignment_target_variables,
is_builtin_variable,
unwrap_parens,
)
from refinery.lib.scripts.ps1.data import PS1_KNOWN_VARIABLES
from refinery.lib.scripts.ps1.deobfuscation.helpers import (
is_array_reverse_call,
iter_variable_mutations,
make_string_literal,
)
from refinery.lib.scripts.ps1.deobfuscation.removal import Ps1RemovalPlans
from refinery.lib.scripts.ps1.deobfuscation.substitution import substitute, substitute_field
from refinery.lib.scripts.ps1.model import (
Ps1ArrayExpression,
Ps1ArrayLiteral,
Ps1AssignmentExpression,
Ps1BinaryExpression,
Ps1CastExpression,
Ps1ClassDefinition,
Ps1DoLoop,
Ps1EnumDefinition,
Ps1ExpressionStatement,
Ps1ForLoop,
Ps1FunctionDefinition,
Ps1HereString,
Ps1IfStatement,
Ps1IndexExpression,
Ps1IntegerLiteral,
Ps1ParenExpression,
Ps1Pipeline,
Ps1PipelineElement,
Ps1RealLiteral,
Ps1ScopeModifier,
Ps1ScriptBlock,
Ps1StringLiteral,
Ps1SwitchStatement,
Ps1TypeExpression,
Ps1UnaryExpression,
Ps1Variable,
Ps1WhileLoop,
)
from refinery.lib.scripts.ps1.synth import Ps1Synthesizer
from refinery.lib.scripts.win32const import DEFAULT_ENVIRONMENT_TEMPLATE
_PS1_DEFAULT_VARIABLES: dict[str, str] = {
key.lower(): value for key, value in {
'ConfirmPreference' : r'High',
'ConsoleFileName' : r'',
'DebugPreference' : r'SilentlyContinue',
'ErrorActionPreference' : r'Continue',
'InformationPreference' : r'SilentlyContinue',
'ProgressPreference' : r'Continue',
'PSCommandPath' : r'',
'PSCulture' : r'en-US',
'PSEmailServer' : r'',
'PSHome' : r'C:\Windows\System32\WindowsPowerShell\v1.0',
'PSScriptRoot' : r'',
'PSSessionApplicationName' : r'wsman',
'PSSessionConfigurationName' : r'http://schemas.microsoft.com/powershell/Microsoft.PowerShell',
'PSUICulture' : r'en-US',
'ShellID' : r'Microsoft.PowerShell',
'VerbosePreference' : r'SilentlyContinue',
'WarningPreference' : r'Continue',
}.items()
}
PS1_ENV_CONSTANTS = {
lower_key: value
for key, value in DEFAULT_ENVIRONMENT_TEMPLATE.items()
if not (lower_key := key.lower()).startswith(('path', 'processor'))
and '{u}' not in value
and '{h}' not in value
}
_PS1_AUTOMATIC_VARIABLES = frozenset({
'_',
'args',
'error',
'event',
'eventargs',
'eventsubscriber',
'executioncontext',
'false',
'foreach',
'home',
'host',
'input',
'lastexitcode',
'matches',
'myinvocation',
'nestedpromptlevel',
'null',
'ofs',
'pid',
'profile',
'psboundparameters',
'pscmdlet',
'pscommandpath',
'psitem',
'psscriptroot',
'psversiontable',
'pwd',
'sender',
'sourceargs',
'sourceeventargs',
'stacktrace',
'switch',
'this',
'true',
})
_PS1_SKIP_VARIABLES = (
_PS1_AUTOMATIC_VARIABLES
| frozenset(PS1_KNOWN_VARIABLES)
| frozenset(_PS1_DEFAULT_VARIABLES)
)
#: The names the engine maintains between statements, so what the script last assigned to one is not
#: what it is worth at the next read: `$_` is rebound per pipeline object, `$Matches` at every
#: `-match`, `$LASTEXITCODE` by every native command, and a preference variable is read by the
#: engine itself. No write of one of these establishes a value this pass may carry to a reader.
_PS1_ENGINE_VARIABLES = _PS1_AUTOMATIC_VARIABLES | frozenset(_PS1_DEFAULT_VARIABLES)
_MIN_EXPANSION_BUDGET = 256
def _collect_mutated_variables(root: Node) -> set[str]:
"""
Return the set of variable keys that are written to anywhere in the AST. This includes
assignment targets, ForEach loop variables, ++/-- operands, and parameter declarations.
"""
mutated: set[str] = set()
for var, _kind, _node in iter_variable_mutations(root):
key = _candidate_key(var)
if key is not None:
mutated.add(key)
for node in root.walk():
if isinstance(node, Ps1ExpressionStatement):
rv = is_array_reverse_call(node)
if rv is not None:
key = _candidate_key(rv)
if key is not None:
mutated.add(key)
return mutated
def _candidate_key(var: Ps1Variable) -> str | None:
"""
Return the candidate lookup key for a variable, or `None` if it is not
eligible for constant inlining.
"""
if var.scope == Ps1ScopeModifier.NONE:
return var.name.lower()
if var.scope == Ps1ScopeModifier.ENV:
return F'env:{var.name.lower()}'
return None
def _constant_value_key(node: Node) -> tuple | None:
"""
Return a hashable key representing the constant value of a node, or `None`
if the node is not constant. Two constant nodes with the same key are
guaranteed to represent the same value.
"""
node = unwrap_parens(node)
if isinstance(node, Ps1IntegerLiteral):
return ('int', node.value)
if isinstance(node, Ps1RealLiteral):
return ('real', node.value)
if isinstance(node, Ps1StringLiteral):
return ('str', node.value)
if isinstance(node, Ps1HereString):
return ('str', node.value)
if isinstance(node, Ps1TypeExpression):
return ('type', node.name)
if is_builtin_variable(node):
return ('var', node.name.lower())
if isinstance(node, Ps1ArrayLiteral):
keys = []
for e in node.elements:
k = _constant_value_key(e)
if k is None:
return None
keys.append(k)
return ('array', tuple(keys))
if isinstance(node, Ps1ArrayExpression):
inner = unwrap_to_array_literal(node)
if inner is not None:
return _constant_value_key(inner)
return None
def _get_array_literal(node: Node) -> Ps1ArrayLiteral | None:
"""
Return the indexable `refinery.lib.scripts.ps1.model.Ps1ArrayLiteral` from either a bare literal
or `@(...)`.
"""
if isinstance(node, Expression):
return unwrap_to_array_literal(node)
return None
def _clone_constant(node: Node) -> Expression:
"""
Create a fresh copy of a constant value node without following parent references. This avoids
the catastrophic cost of `copy.deepcopy` which traverses the entire AST through parents.
"""
unwrapped = unwrap_parens(node)
if isinstance(unwrapped, Ps1ArrayExpression):
inner = unwrap_to_array_literal(unwrapped)
if inner is None:
raise TypeError(F'cannot clone {type(unwrapped).__name__}')
unwrapped = inner
if not isinstance(unwrapped, Expression):
raise TypeError(F'cannot clone {type(unwrapped).__name__}')
clone = _clone_node(unwrapped)
if isinstance(clone, Ps1ArrayLiteral) and len(clone.elements) > 1:
return Ps1ParenExpression(expression=clone)
return clone
def _walk_outer_scope(root: Node):
"""
Walk the AST like `root.walk()` but skip the bodies of function, class, and enum definitions.
The definition node itself is yielded so that it can still be removed or inspected.
"""
stack: list[Node] = [root]
while stack:
node = stack.pop()
yield node
if isinstance(node, (Ps1FunctionDefinition, Ps1ClassDefinition, Ps1EnumDefinition)):
continue
for child in node.children():
stack.append(child)
def _find_removable_statement(node: Node) -> Node | None:
"""
Walk upward from an expression node to find the statement-level node that can be removed from
its parent's body list.
"""
cursor = node
while cursor.parent is not None:
parent = cursor.parent
if isinstance(parent, Ps1ExpressionStatement):
cursor = parent
continue
if isinstance(parent, Ps1PipelineElement):
cursor = parent
continue
if isinstance(parent, Ps1Pipeline):
if len(parent.elements) == 1:
cursor = parent
continue
return cursor
return None
class _ConstantTable:
"""
The constant value each write of a script establishes, keyed by the identity of the occurrence
that writes it, and the constants of the names the script never writes at all.
Two tables because they answer two questions. A write is a point in the program and the flow
model orders it against a read; an ambient constant is a value the engine established before the
script ran — `$env:ComSpec`, `$ErrorActionPreference` — and there is no point to order it
against. The pass this replaces held one table for both and used the absence of a position as
the marker, so a write it could not place and a value that has no position were the same entry.
A write with no entry here is not a lesser kind of write. It is a write whose value this pass has
nothing to say about, and the flow model orders and kills it exactly as it does any other:
sorting writes by whether their value happens to be constant is what made `if ($c) { $x = 'b' }`
fold and `if ($c) { $x = $y }` refuse.
"""
def __init__(self, root: Node, flow: Ps1VariableFlow):
self.by_write: dict[int, Expression] = {}
self.ambient: dict[str, Expression] = {}
self.values: defaultdict[str, list[Expression]] = defaultdict(list)
self._collect_writes(root)
self._collect_ambient(root, flow)
def _collect_writes(self, root: Node):
for node in root.walk():
if not isinstance(node, Ps1AssignmentExpression):
continue
if node.operator != '=' or node.value is None:
continue
targets = assignment_target_variables(node.target)
if len(targets) != 1:
continue
key = _candidate_key(targets[0])
if key is None or key in _PS1_ENGINE_VARIABLES:
# A preference or automatic variable is the engine's as much as the script's: it
# reads and writes these names between statements — `$_` per pipeline object,
# `$Matches` at every `-match` — so what the script last assigned is not what the
# name is worth. The ambient table is the only thing that may answer for one, and it
# answers only while the script leaves the name alone.
continue
if _constant_value_key(node.value) is None:
continue
value = unwrap_parens(node.value)
self.by_write[id(targets[0])] = value
self.values[key].append(value)
def _collect_ambient(self, root: Node, flow: Ps1VariableFlow):
"""
A default the engine supplies is only this name's value while the script leaves the name
alone. Any write of it anywhere, and any statement that reaches into its value, replaces the
default with something this table has no claim on — including a write inside a block, which
`Ps1SemanticModel` binds locally but `. { }` performs on the caller.
A write nobody can attribute is *not* collected here, because it is not a fact about the
whole script: it lands at a point, and an ambient default is a definition at the script's
entry, so the two are ordered like anything else. `Ps1VariableFlow.ambient_value_survives`
asks that per read — silencing every default here instead was measured, and it costs the
`$PSHome` unpacking that an obfuscated loader's first stage is built out of.
"""
touched = set(flow.mutated_in_place)
for node in root.walk():
if isinstance(node, Ps1Variable) and is_write_occurrence(node):
touched.add(binding_key(node))
for key, value in _PS1_DEFAULT_VARIABLES.items():
if key not in touched:
self._add_ambient(key, make_string_literal(value))
for name, value in PS1_ENV_CONSTANTS.items():
key = F'env:{name}'
if key not in touched:
self._add_ambient(key, make_string_literal(value))
def _add_ambient(self, key: str, value: Expression):
self.ambient[key] = value
self.values[key].append(value)
def __bool__(self) -> bool:
return bool(self.by_write or self.ambient)
class _InlineRecord:
"""
The read occurrences one substitution walk replaced, per binding.
The count is what licenses removing the binding's writes — every occurrence in `Binding.reads`
has to be accounted for, and that set includes reads this pass never walked, such as one inside a
function body — and the replacement nodes are where the value now stands, so they are what the
pass can point at when it claims that removing the write destroys nothing.
"""
def __init__(self):
self._bindings: dict[int, Binding] = {}
self._replacements: defaultdict[int, list[Node]] = defaultdict(list)
def add(self, binding: Binding, replacement: Node):
self._bindings[id(binding)] = binding
self._replacements[id(binding)].append(replacement)
def __iter__(self) -> Iterator[tuple[Binding, list[Node]]]:
for key, binding in self._bindings.items():
yield binding, self._replacements[key]
class _Inlining:
"""
The state one substitution walk carries: the constants it may install, the flow model that says
which of them a read observes, the keys the expansion budget has already refused, the variable
occurrences an enclosing index expression has already spoken for, and what was replaced.
"""
def __init__(self, table: _ConstantTable, flow: Ps1VariableFlow, blocked: set[str]):
self.table = table
self.flow = flow
self.blocked = blocked
self.handled: set[int] = set()
self.record = _InlineRecord()
def value_at(self, var: Ps1Variable, key: str) -> Expression | None:
"""
The constant *var* holds where it stands, or `None` when no single value does.
"""
binding = self.binding_of(var)
if binding is None:
value = self.table.ambient.get(key)
if value is None or not self.flow.ambient_value_survives(var):
return None
return value
write = self.flow.reaching_definition(var)
if write is None:
return None
return self.table.by_write.get(id(write))
def binding_of(self, var: Ps1Variable) -> Binding | None:
return self.flow.semantic.binding_of(var)
def installed(self, var: Ps1Variable, replacement: Node):
binding = self.binding_of(var)
if binding is not None:
self.record.add(binding, replacement)
class Ps1ConstantInlining(Transformer):
def __init__(self, max_expansion_ratio: float = 0.2, min_inlines_to_prune: int | None = 1):
super().__init__()
self.max_expansion_ratio = max_expansion_ratio
self.min_inlines_to_prune = min_inlines_to_prune
def visit(self, node: Node):
# Captured once rather than re-read per reference: every substitution below marks the pass
# changed, which drops the cache, so a per-site lookup would rebuild the control-flow graphs
# of the whole script once per inlined variable. Nothing this pass adds or removes is a
# statement, so the graphs it would rebuild are the graphs it already has.
flow = model_cache(self, node).variable_flow
table = _ConstantTable(node, flow)
if not table:
return None
state = _Inlining(table, flow, self._blocked_by_expansion(node, table))
self._substitute(node, state)
self._remove_dead_assignments(table, state)
return None
def _blocked_by_expansion(self, root: Node, table: _ConstantTable) -> set[str]:
"""
The keys whose substitution would grow the script past the expansion budget, estimated over
every reference before any of them is installed. Purely a size heuristic: it withholds an
inlining that is correct, and it is asked before the flow model so that a script full of
references to one large array does not pay for a reaching-definition query per reference.
"""
synth = Ps1Synthesizer()
script_size = len(synth.convert(root))
max_budget = max(_MIN_EXPANSION_BUDGET, int(script_size * self.max_expansion_ratio))
value_lengths: dict[str, int] = {}
array_literals: dict[str, Ps1ArrayLiteral | None] = {}
for key, values in table.values.items():
value_lengths[key] = max(len(synth.convert(value)) for value in values)
array_literals[key] = _get_array_literal(values[0])
elem_lengths: dict[tuple[str, int], int] = {}
expansion: defaultdict[str, int] = defaultdict(int)
for node in _walk_outer_scope(root):
if isinstance(node, Ps1IndexExpression):
var = node.object
if not isinstance(var, Ps1Variable):
continue
key = _candidate_key(var)
if key is None or key not in table.values or node.index is None:
continue
if isinstance(node.index, Ps1IntegerLiteral):
array = array_literals[key]
if array is None:
continue
idx = node.index.value
if not 0 <= idx < len(array.elements):
continue
ref_len = 1 + len(var.name) + 1 + len(node.index.raw) + 1
cache_key = (key, idx)
if cache_key not in elem_lengths:
elem_lengths[cache_key] = len(synth.convert(array.elements[idx]))
expansion[key] += max(0, elem_lengths[cache_key] - ref_len)
elif isinstance(table.values[key][0], (Ps1StringLiteral, Ps1HereString)):
expansion[key] += max(0, value_lengths[key] - (1 + len(var.name)))
elif isinstance(node, Ps1Variable):
key = _candidate_key(node)
if key is None or key not in table.values or is_write_occurrence(node):
continue
expansion[key] += max(0, value_lengths[key] - (1 + len(node.name)))
return {key for key in table.values if expansion[key] > max_budget}
def _substitute(self, root: Node, state: _Inlining):
"""
Replace every reference this pass can resolve with the value it observes.
The walk stops at a function, class, or enum body. What a read inside one observes is a
question about the call sites that reach it, which the flow model refuses rather than
answers, so descending would only spend a query per reference to be told nothing; and a
class body is opaque to the graphs, so a property initializer inside one locates to the
class statement and would be ordered against code it does not run beside.
Which positions may hold a value at all is
`refinery.lib.scripts.ps1.analysis.model.is_substitutable_position`, asked once here rather
than reassembled from the positional predicates it is made of. It is a fact about the
*position*, not about the binding: the flow model refuses a name stored through as well,
via `Ps1FlowUnknown.MUTATED_IN_PLACE`, but the ambient table answers with no binding at all
and would otherwise install a constant where `$x[0] = 'z'` names a place rather than a
value.
"""
for node in list(_walk_outer_scope(root)):
if isinstance(node, Ps1IndexExpression):
var = node.object
if not isinstance(var, Ps1Variable):
continue
# Spoken for either way: the walk snapshot still holds this occurrence after the
# index expression around it has been swapped out, and substituting it a second time
# would install the whole value where an element of it now stands.
state.handled.add(id(var))
key = _candidate_key(var)
if key is not None and is_substitutable_position(var):
self._substitute_index_reference(node, var, key, state)
elif isinstance(node, Ps1Variable):
if id(node) in state.handled or not is_substitutable_position(node):
continue
key = _candidate_key(node)
if key is not None and key not in state.blocked:
self._substitute_variable_reference(node, key, state)
def _substitute_index_reference(
self,
node: Ps1IndexExpression,
var: Ps1Variable,
key: str,
state: _Inlining,
) -> None:
const_value = state.value_at(var, key)
if const_value is None:
return
if not isinstance(node.index, Ps1IntegerLiteral):
if key in state.blocked or not isinstance(const_value, Ps1StringLiteral):
return
replacement = _clone_constant(const_value)
if substitute_field(node, 'object', replacement):
self.mark_changed()
state.installed(var, replacement)
return
idx = node.index.value
if isinstance(const_value, Ps1StringLiteral):
text = const_value.value
if not 0 <= idx < len(text):
return
replacement = make_string_literal(text[idx])
else:
array = _get_array_literal(const_value)
if array is None or not 0 <= idx < len(array.elements):
return
replacement = _clone_constant(array.elements[idx])
if substitute(node, replacement):
self.mark_changed()
state.installed(var, replacement)
def _substitute_variable_reference(
self,
node: Ps1Variable,
key: str,
state: _Inlining,
) -> None:
const_value = state.value_at(node, key)
if const_value is None:
return
replacement = _clone_constant(const_value)
if substitute(node, replacement):
self.mark_changed()
state.installed(node, replacement)
def _remove_dead_assignments(self, table: _ConstantTable, state: _Inlining):
"""
Delete the constant writes of every binding whose value nothing observes any more.
Removal is decided per binding, and counted against that binding's own reads rather than
against the references this walk resolved. Every occurrence in `Binding.reads` observes the
value, including the ones the walk cannot answer for — a read inside a function body, a read
after a write whose value is not constant, an index this pass cannot evaluate — and each of
them is a reader the write still has. Counting only what the walk replaced is what let
`$x = 'a'; function f { Write-Host $x }; Write-Host $x; f` delete the assignment `f` reads.
A write that observes the previous value is a read as much as a write, so a binding with one
is never dead however many of its reads were substituted; and a write whose value this pass
holds no constant for stays, because deleting it would drop whatever it does to produce that
value.
`Binding.reads` is the whole list of readers only when the binding ends with its own body.
`refinery.lib.scripts.ps1.analysis.model.Ps1SemanticModel` binds a bare write to the
block it is written in, and `refinery.lib.scripts.ps1.analysis.blocks.Ps1BlockModel` is
the layer that says a `. { }` or a `ForEach-Object` body performs that write on whoever runs
it — where the readers are other bindings entirely, and `_block_kills` already honours the
same fact on the read side. Deleting such a write leaves the caller reading the value from
before the body.
"""
plans = Ps1RemovalPlans()
for binding, replacements in state.record:
if len(replacements) < len(binding.reads):
continue
if self._writes_leave_the_body(state.flow, binding):
continue
if any(write.role is Ps1OccurrenceRole.WRITE_OBSERVING for write in binding.writes):
continue
if self.min_inlines_to_prune is not None:
if len(replacements) < self.min_inlines_to_prune:
continue
for write in binding.writes:
assignment = assignment_of(write.node)
if assignment is None or id(write.node) not in table.by_write:
continue
statement = self._find_removable_statement(assignment)
if statement is not None:
plans.propose(statement)
if plans.commit():
self.mark_changed()
@staticmethod
def _writes_leave_the_body(flow: Ps1VariableFlow, binding: Binding) -> bool:
"""
Whether *binding*'s writes land in the scope of whatever runs the body they are written in,
rather than in a scope that ends with that body. True for every block but a proven child
scope — see `refinery.lib.scripts.ps1.analysis.blocks` for why that asymmetry is the safe
one.
"""
node = binding.scope.node
return isinstance(node, Ps1ScriptBlock) and flow.blocks.may_write_caller_scope(node)
_find_removable_statement = staticmethod(_find_removable_statement)
class Ps1NullVariableInlining(Transformer):
"""
Replace references to never-assigned variables with `$Null`. Only operates on variables that
appear in expression contexts where null coercion enables further simplification (arithmetic,
comparison, cast, assignment value).
"""
@staticmethod
def _is_null_eligible(ref: Ps1Variable) -> bool:
cursor = ref
while cursor.parent is not None:
parent = cursor.parent
if isinstance(parent, Ps1BinaryExpression):
return True
if isinstance(parent, Ps1UnaryExpression):
return True
if isinstance(parent, Ps1CastExpression):
cursor = parent
continue
if isinstance(parent, Ps1AssignmentExpression) and cursor is parent.value:
return True
if isinstance(parent, (Ps1ParenExpression, Ps1ArrayLiteral)):
cursor = parent
continue
if isinstance(parent, (Ps1WhileLoop, Ps1DoLoop, Ps1ForLoop)) and cursor is parent.condition:
return True
if isinstance(parent, (Ps1IfStatement, Ps1SwitchStatement)):
return any(cursor is cond for cond, _ in parent.clauses)
return False
return False
def visit(self, node: Node):
mutated = _collect_mutated_variables(node)
for ref in list(node.walk()):
if not isinstance(ref, Ps1Variable):
continue
key = _candidate_key(ref)
if key is None:
continue
if key in mutated:
continue
if key in PS1_KNOWN_VARIABLES:
continue
if key in _PS1_DEFAULT_VARIABLES:
continue
if key in _PS1_AUTOMATIC_VARIABLES:
continue
if key.startswith('env:'):
continue
if is_assignment_write_target(ref):
continue
if not self._is_null_eligible(ref):
continue
if not substitute(ref, Ps1Variable(name='Null')):
continue
self.mark_changed()
Classes
class Ps1ConstantInlining (max_expansion_ratio=0.2, min_inlines_to_prune=1)-
In-place tree rewriter. Each visit method may return a replacement node or
Noneto keep the original. Tracks whether any transformation was applied via thechangedflag.When a
modelscache is attached by the pipeline, settingchangedtruthy invalidates it, so a transform that mutates the tree never leaves a stale model behind for the next consumer.Expand source code Browse git
class Ps1ConstantInlining(Transformer): def __init__(self, max_expansion_ratio: float = 0.2, min_inlines_to_prune: int | None = 1): super().__init__() self.max_expansion_ratio = max_expansion_ratio self.min_inlines_to_prune = min_inlines_to_prune def visit(self, node: Node): # Captured once rather than re-read per reference: every substitution below marks the pass # changed, which drops the cache, so a per-site lookup would rebuild the control-flow graphs # of the whole script once per inlined variable. Nothing this pass adds or removes is a # statement, so the graphs it would rebuild are the graphs it already has. flow = model_cache(self, node).variable_flow table = _ConstantTable(node, flow) if not table: return None state = _Inlining(table, flow, self._blocked_by_expansion(node, table)) self._substitute(node, state) self._remove_dead_assignments(table, state) return None def _blocked_by_expansion(self, root: Node, table: _ConstantTable) -> set[str]: """ The keys whose substitution would grow the script past the expansion budget, estimated over every reference before any of them is installed. Purely a size heuristic: it withholds an inlining that is correct, and it is asked before the flow model so that a script full of references to one large array does not pay for a reaching-definition query per reference. """ synth = Ps1Synthesizer() script_size = len(synth.convert(root)) max_budget = max(_MIN_EXPANSION_BUDGET, int(script_size * self.max_expansion_ratio)) value_lengths: dict[str, int] = {} array_literals: dict[str, Ps1ArrayLiteral | None] = {} for key, values in table.values.items(): value_lengths[key] = max(len(synth.convert(value)) for value in values) array_literals[key] = _get_array_literal(values[0]) elem_lengths: dict[tuple[str, int], int] = {} expansion: defaultdict[str, int] = defaultdict(int) for node in _walk_outer_scope(root): if isinstance(node, Ps1IndexExpression): var = node.object if not isinstance(var, Ps1Variable): continue key = _candidate_key(var) if key is None or key not in table.values or node.index is None: continue if isinstance(node.index, Ps1IntegerLiteral): array = array_literals[key] if array is None: continue idx = node.index.value if not 0 <= idx < len(array.elements): continue ref_len = 1 + len(var.name) + 1 + len(node.index.raw) + 1 cache_key = (key, idx) if cache_key not in elem_lengths: elem_lengths[cache_key] = len(synth.convert(array.elements[idx])) expansion[key] += max(0, elem_lengths[cache_key] - ref_len) elif isinstance(table.values[key][0], (Ps1StringLiteral, Ps1HereString)): expansion[key] += max(0, value_lengths[key] - (1 + len(var.name))) elif isinstance(node, Ps1Variable): key = _candidate_key(node) if key is None or key not in table.values or is_write_occurrence(node): continue expansion[key] += max(0, value_lengths[key] - (1 + len(node.name))) return {key for key in table.values if expansion[key] > max_budget} def _substitute(self, root: Node, state: _Inlining): """ Replace every reference this pass can resolve with the value it observes. The walk stops at a function, class, or enum body. What a read inside one observes is a question about the call sites that reach it, which the flow model refuses rather than answers, so descending would only spend a query per reference to be told nothing; and a class body is opaque to the graphs, so a property initializer inside one locates to the class statement and would be ordered against code it does not run beside. Which positions may hold a value at all is `refinery.lib.scripts.ps1.analysis.model.is_substitutable_position`, asked once here rather than reassembled from the positional predicates it is made of. It is a fact about the *position*, not about the binding: the flow model refuses a name stored through as well, via `Ps1FlowUnknown.MUTATED_IN_PLACE`, but the ambient table answers with no binding at all and would otherwise install a constant where `$x[0] = 'z'` names a place rather than a value. """ for node in list(_walk_outer_scope(root)): if isinstance(node, Ps1IndexExpression): var = node.object if not isinstance(var, Ps1Variable): continue # Spoken for either way: the walk snapshot still holds this occurrence after the # index expression around it has been swapped out, and substituting it a second time # would install the whole value where an element of it now stands. state.handled.add(id(var)) key = _candidate_key(var) if key is not None and is_substitutable_position(var): self._substitute_index_reference(node, var, key, state) elif isinstance(node, Ps1Variable): if id(node) in state.handled or not is_substitutable_position(node): continue key = _candidate_key(node) if key is not None and key not in state.blocked: self._substitute_variable_reference(node, key, state) def _substitute_index_reference( self, node: Ps1IndexExpression, var: Ps1Variable, key: str, state: _Inlining, ) -> None: const_value = state.value_at(var, key) if const_value is None: return if not isinstance(node.index, Ps1IntegerLiteral): if key in state.blocked or not isinstance(const_value, Ps1StringLiteral): return replacement = _clone_constant(const_value) if substitute_field(node, 'object', replacement): self.mark_changed() state.installed(var, replacement) return idx = node.index.value if isinstance(const_value, Ps1StringLiteral): text = const_value.value if not 0 <= idx < len(text): return replacement = make_string_literal(text[idx]) else: array = _get_array_literal(const_value) if array is None or not 0 <= idx < len(array.elements): return replacement = _clone_constant(array.elements[idx]) if substitute(node, replacement): self.mark_changed() state.installed(var, replacement) def _substitute_variable_reference( self, node: Ps1Variable, key: str, state: _Inlining, ) -> None: const_value = state.value_at(node, key) if const_value is None: return replacement = _clone_constant(const_value) if substitute(node, replacement): self.mark_changed() state.installed(node, replacement) def _remove_dead_assignments(self, table: _ConstantTable, state: _Inlining): """ Delete the constant writes of every binding whose value nothing observes any more. Removal is decided per binding, and counted against that binding's own reads rather than against the references this walk resolved. Every occurrence in `Binding.reads` observes the value, including the ones the walk cannot answer for — a read inside a function body, a read after a write whose value is not constant, an index this pass cannot evaluate — and each of them is a reader the write still has. Counting only what the walk replaced is what let `$x = 'a'; function f { Write-Host $x }; Write-Host $x; f` delete the assignment `f` reads. A write that observes the previous value is a read as much as a write, so a binding with one is never dead however many of its reads were substituted; and a write whose value this pass holds no constant for stays, because deleting it would drop whatever it does to produce that value. `Binding.reads` is the whole list of readers only when the binding ends with its own body. `refinery.lib.scripts.ps1.analysis.model.Ps1SemanticModel` binds a bare write to the block it is written in, and `refinery.lib.scripts.ps1.analysis.blocks.Ps1BlockModel` is the layer that says a `. { }` or a `ForEach-Object` body performs that write on whoever runs it — where the readers are other bindings entirely, and `_block_kills` already honours the same fact on the read side. Deleting such a write leaves the caller reading the value from before the body. """ plans = Ps1RemovalPlans() for binding, replacements in state.record: if len(replacements) < len(binding.reads): continue if self._writes_leave_the_body(state.flow, binding): continue if any(write.role is Ps1OccurrenceRole.WRITE_OBSERVING for write in binding.writes): continue if self.min_inlines_to_prune is not None: if len(replacements) < self.min_inlines_to_prune: continue for write in binding.writes: assignment = assignment_of(write.node) if assignment is None or id(write.node) not in table.by_write: continue statement = self._find_removable_statement(assignment) if statement is not None: plans.propose(statement) if plans.commit(): self.mark_changed() @staticmethod def _writes_leave_the_body(flow: Ps1VariableFlow, binding: Binding) -> bool: """ Whether *binding*'s writes land in the scope of whatever runs the body they are written in, rather than in a scope that ends with that body. True for every block but a proven child scope — see `refinery.lib.scripts.ps1.analysis.blocks` for why that asymmetry is the safe one. """ node = binding.scope.node return isinstance(node, Ps1ScriptBlock) and flow.blocks.may_write_caller_scope(node) _find_removable_statement = staticmethod(_find_removable_statement)Ancestors
Methods
def visit(self, node)-
Expand source code Browse git
def visit(self, node: Node): # Captured once rather than re-read per reference: every substitution below marks the pass # changed, which drops the cache, so a per-site lookup would rebuild the control-flow graphs # of the whole script once per inlined variable. Nothing this pass adds or removes is a # statement, so the graphs it would rebuild are the graphs it already has. flow = model_cache(self, node).variable_flow table = _ConstantTable(node, flow) if not table: return None state = _Inlining(table, flow, self._blocked_by_expansion(node, table)) self._substitute(node, state) self._remove_dead_assignments(table, state) return None
Inherited members
class Ps1NullVariableInlining-
Replace references to never-assigned variables with
$Null. Only operates on variables that appear in expression contexts where null coercion enables further simplification (arithmetic, comparison, cast, assignment value).Expand source code Browse git
class Ps1NullVariableInlining(Transformer): """ Replace references to never-assigned variables with `$Null`. Only operates on variables that appear in expression contexts where null coercion enables further simplification (arithmetic, comparison, cast, assignment value). """ @staticmethod def _is_null_eligible(ref: Ps1Variable) -> bool: cursor = ref while cursor.parent is not None: parent = cursor.parent if isinstance(parent, Ps1BinaryExpression): return True if isinstance(parent, Ps1UnaryExpression): return True if isinstance(parent, Ps1CastExpression): cursor = parent continue if isinstance(parent, Ps1AssignmentExpression) and cursor is parent.value: return True if isinstance(parent, (Ps1ParenExpression, Ps1ArrayLiteral)): cursor = parent continue if isinstance(parent, (Ps1WhileLoop, Ps1DoLoop, Ps1ForLoop)) and cursor is parent.condition: return True if isinstance(parent, (Ps1IfStatement, Ps1SwitchStatement)): return any(cursor is cond for cond, _ in parent.clauses) return False return False def visit(self, node: Node): mutated = _collect_mutated_variables(node) for ref in list(node.walk()): if not isinstance(ref, Ps1Variable): continue key = _candidate_key(ref) if key is None: continue if key in mutated: continue if key in PS1_KNOWN_VARIABLES: continue if key in _PS1_DEFAULT_VARIABLES: continue if key in _PS1_AUTOMATIC_VARIABLES: continue if key.startswith('env:'): continue if is_assignment_write_target(ref): continue if not self._is_null_eligible(ref): continue if not substitute(ref, Ps1Variable(name='Null')): continue self.mark_changed()Ancestors
Methods
def visit(self, node)-
Expand source code Browse git
def visit(self, node: Node): mutated = _collect_mutated_variables(node) for ref in list(node.walk()): if not isinstance(ref, Ps1Variable): continue key = _candidate_key(ref) if key is None: continue if key in mutated: continue if key in PS1_KNOWN_VARIABLES: continue if key in _PS1_DEFAULT_VARIABLES: continue if key in _PS1_AUTOMATIC_VARIABLES: continue if key.startswith('env:'): continue if is_assignment_write_target(ref): continue if not self._is_null_eligible(ref): continue if not substitute(ref, Ps1Variable(name='Null')): continue self.mark_changed()
Inherited members