Module refinery.lib.scripts.js.deobfuscation.dispatcher
The dispatcher obfuscation wraps function bodies into a central routing function that uses a string keyed lookup table and a global payload array for argument passing. This transformer detects the pattern structurally (no reliance on variable names), extracts the original functions, rewrites all call sites, and removes the dispatcher scaffolding.
Expand source code Browse git
"""
The dispatcher obfuscation wraps function bodies into a central routing function that uses a string
keyed lookup table and a global payload array for argument passing. This transformer detects the
pattern structurally (no reliance on variable names), extracts the original functions, rewrites all
call sites, and removes the dispatcher scaffolding.
"""
from __future__ import annotations
from dataclasses import dataclass
from refinery.lib.scripts import (
Node,
_clone_node,
_remove_from_parent,
_replace_in_parent,
)
from refinery.lib.scripts.js.analysis.cache import model_cache
from refinery.lib.scripts.js.analysis.model import enclosing_operator
from refinery.lib.scripts.js.deobfuscation.helpers import (
ScopeProcessingTransformer,
access_key,
binding_has_references,
make_undefined_expression,
property_key,
remove_declarator,
)
from refinery.lib.scripts.js.model import (
JsArrayExpression,
JsArrayPattern,
JsAssignmentExpression,
JsBinaryExpression,
JsBlockStatement,
JsCallExpression,
JsExpressionStatement,
JsFunctionDeclaration,
JsFunctionExpression,
JsIdentifier,
JsIfStatement,
JsLogicalExpression,
JsMemberExpression,
JsNewExpression,
JsNullLiteral,
JsObjectExpression,
JsProperty,
JsReturnStatement,
JsScript,
JsSequenceExpression,
JsStringLiteral,
JsVariableDeclaration,
JsVariableDeclarator,
is_async_function,
is_generator_function,
strip_parens,
wraps_return,
)
class _Unreadable:
"""
What a dispatch argument this pass cannot read as a fixed string is reported as, kept apart from
the `None` an absent argument gives: an argument that is not there is one the dispatcher sees as
`undefined` and compares unequal to every flag, while one that is there and unreadable may be
any of them and may run something on the way.
"""
_UNREADABLE = _Unreadable()
def _flag_argument(
call: JsCallExpression | JsNewExpression,
index: int,
) -> str | None | _Unreadable:
"""
The string the argument of *call* at *index* certainly evaluates to, `None` where the call has
no such argument, and `_UNREADABLE` where it has one this pass cannot read as a fixed string.
"""
if len(call.arguments) <= index:
return None
argument = strip_parens(call.arguments[index])
if isinstance(argument, JsStringLiteral):
return argument.value
return _UNREADABLE
def _reads_the_payload(fn: JsFunctionExpression, payload_id: str) -> bool:
"""
Whether the table entry *fn* takes its arguments out of the payload, so what the payload holds
when it is reached decides what it computes.
"""
return _extract_params(fn, payload_id) != []
@dataclass
class _DispatcherInfo:
"""
All structurally-extracted metadata about a single dispatcher function.
The three flag strings are carried beside the things they select, because what a dispatch means
is decided by the arguments after the key and not by the key alone. *init_flag* is the value of
the second parameter that empties the payload, so a dispatch passing it reaches its callee with
no arguments however the payload was filled. *create_flag* is the value that makes the
dispatcher hand back the table entry instead of calling it, so a dispatch passing it names a
function where every other one names a result. *wrap_flag* is the value of the third parameter
that wraps the result in an object under `wrap_key`, so a dispatch passing it denotes that
object and only the access on that key denotes the result.
"""
decl: JsFunctionDeclaration
dispatcher_id: str
fns_map: dict[str, JsFunctionExpression]
fns_declarator: JsVariableDeclarator
payload_id: str
init_flag: str
wrap_key: str | None
wrap_flag: str | None
cache_id: str | None
create_flag: str | None
@dataclass
class _DispatchSite:
"""
One dispatch this pass can read: the expression it is written as, the identifier naming the
dispatcher within it, the name of the function it selects, and the arguments that function is
reached with. *arguments* is `None` where the site names the function rather than calling it,
which the wrapped-reference form does.
*reference* is the one occurrence of the dispatcher name this site consumes, and it is what the
coverage question is asked over. Asking it over the site's whole subtree instead would count a
dispatch this pass cannot read as covered merely for standing inside one it can, and a payload
argument is carried into the replacement rather than discarded with the rest of it.
"""
node: Node
reference: JsIdentifier
key: str
arguments: list | None
def replacement(self) -> Node:
if self.arguments is None:
return JsIdentifier(name=self.key)
return JsCallExpression(
callee=JsIdentifier(name=self.key),
arguments=self.arguments,
)
def _extract_fns_table(
body: list,
) -> tuple[JsVariableDeclarator, dict[str, JsFunctionExpression]] | None:
"""
Finds a declaration of the form
var fns = { ... }
where every property value is a zero-parameter
`refinery.lib.scripts.js.model.JsFunctionExpression`. Returns the declarator node and a map from
string key to function.
"""
for stmt in body:
if not isinstance(stmt, JsVariableDeclaration):
continue
for decl in stmt.declarations:
if not isinstance(decl, JsVariableDeclarator):
continue
if not isinstance(decl.init, JsObjectExpression):
continue
obj = decl.init
if not obj.properties:
continue
fns: dict[str, JsFunctionExpression] = {}
ok = True
for prop in obj.properties:
if not isinstance(prop, JsProperty):
ok = False
break
key = property_key(prop)
if key is None:
ok = False
break
if not isinstance(prop.value, JsFunctionExpression):
ok = False
break
if prop.value.params:
ok = False
break
fns[key] = prop.value
if ok and fns:
return decl, fns
return None
def _guarded_flag(stmt: Node, param: str) -> str | None:
"""
The string *stmt*'s guard compares *param* to, or `None` where *stmt* is not an `if` guarded
that way. Every branch this pass reads out of a dispatcher body is selected by one such
comparison, and the value compared against is what a dispatch has to spell to take the branch —
so the same read serves finding the branch and deciding whether a call site enters it.
"""
if not isinstance(stmt, JsIfStatement):
return None
test = stmt.test
if not isinstance(test, JsBinaryExpression) or test.operator != '===':
return None
if not isinstance(test.left, JsIdentifier) or test.left.name != param:
return None
if not isinstance(test.right, JsStringLiteral):
return None
return test.right.value
def _find_payload_id(body: list, second_param: str) -> tuple[str, str] | None:
"""
Find the payload-init guard:
if (p1 === "...") { payload = []; }
and return the payload identifier name together with the flag that empties it. The guard
compares the function's second parameter to a string literal and assigns an empty array to the
payload variable.
"""
for stmt in body:
flag = _guarded_flag(stmt, second_param)
if flag is None:
continue
assert isinstance(stmt, JsIfStatement)
cons = stmt.consequent
if isinstance(cons, JsBlockStatement) and len(cons.body) == 1:
cons = cons.body[0]
if not isinstance(cons, JsExpressionStatement):
continue
expr = cons.expression
if not isinstance(expr, JsAssignmentExpression) or expr.operator != '=':
continue
if isinstance(expr.left, JsIdentifier) and isinstance(expr.right, JsArrayExpression):
if not expr.right.elements:
return expr.left.name, flag
return None
def _find_wrap_key(body: list, third_param: str) -> tuple[str, str] | None:
"""
Find the return-type wrapper:
if (p2 === "...") { return { "wrapKey": output }; }
and return the wrapper property name together with the flag that asks for it.
"""
for stmt in body:
flag = _guarded_flag(stmt, third_param)
if flag is None:
continue
assert isinstance(stmt, JsIfStatement)
cons = stmt.consequent
if isinstance(cons, JsBlockStatement) and len(cons.body) == 1:
inner = cons.body[0]
else:
inner = cons
if not isinstance(inner, JsReturnStatement):
continue
ret_val = inner.argument
if not isinstance(ret_val, JsObjectExpression):
continue
if len(ret_val.properties) != 1:
continue
prop = ret_val.properties[0]
if isinstance(prop, JsProperty):
key = property_key(prop)
if key is not None:
return key, flag
return None
def _find_cache_id(body: list, first_param: str, second_param: str) -> tuple[str, str] | None:
"""
Find the cache variable from the create-flag branch. Looks for an `if` guarded on the second
parameter whose body contains a logical-or assignment like
cache[p0] || (cache[p0] = ...)
Returns the cache identifier together with the flag that reaches that branch. The guard is read
rather than skipped past, because the branch does not call the entry it looks up: a dispatch
that spells this flag names the function, and one that does not names what calling it returned.
"""
for stmt in body:
flag = _guarded_flag(stmt, second_param)
if flag is None:
continue
for node in stmt.walk():
if not isinstance(node, JsMemberExpression):
continue
if (
isinstance(node.object, JsIdentifier)
and isinstance(node.property, JsIdentifier)
and node.property.name == first_param
and node.computed
):
parent = node.parent
if isinstance(parent, JsLogicalExpression) and parent.operator == '||':
return node.object.name, flag
return None
def _detect_dispatcher(func: JsFunctionDeclaration) -> _DispatcherInfo | None:
"""
Structurally detect whether `func` is a dispatcher function. Returns the extracted metadata or
`None` if the function does not match the pattern.
"""
if wraps_return(func):
return None
if not isinstance(func.id, JsIdentifier):
return None
if not isinstance(func.body, JsBlockStatement):
return None
if len(func.params) < 3:
return None
p0 = func.params[0]
p1 = func.params[1]
p2 = func.params[2]
if (
not isinstance(p0, JsIdentifier)
or not isinstance(p1, JsIdentifier)
or not isinstance(p2, JsIdentifier)
):
return None
first_param: str = p0.name
second_param: str = p1.name
third_param: str = p2.name
body = func.body.body
result = _extract_fns_table(body)
if result is None:
return None
fns_declarator, fns_map = result
payload = _find_payload_id(body, second_param)
if payload is None:
return None
payload_id, init_flag = payload
wrap = _find_wrap_key(body, third_param)
cache = _find_cache_id(body, first_param, second_param)
return _DispatcherInfo(
decl=func,
dispatcher_id=func.id.name,
fns_map=fns_map,
fns_declarator=fns_declarator,
payload_id=payload_id,
init_flag=init_flag,
wrap_key=None if wrap is None else wrap[0],
wrap_flag=None if wrap is None else wrap[1],
cache_id=None if cache is None else cache[0],
create_flag=None if cache is None else cache[1],
)
def _extract_params(
fn: JsFunctionExpression,
payload_id: str,
) -> list[JsIdentifier] | None:
"""
Extract parameter names from the leading payload destructuring statement:
var [a, b] = payload;
Returns the parameter identifiers or `None` if the pattern is not found.
"""
if not isinstance(fn.body, JsBlockStatement) or not fn.body.body:
return []
first = fn.body.body[0]
if not isinstance(first, JsVariableDeclaration):
return []
for decl in first.declarations:
if not isinstance(decl, JsVariableDeclarator):
continue
if not isinstance(decl.id, JsArrayPattern):
continue
if not isinstance(decl.init, JsIdentifier):
continue
if decl.init.name != payload_id:
continue
params: list[JsIdentifier] = []
for elem in decl.id.elements:
if not isinstance(elem, JsIdentifier):
return None
params.append(JsIdentifier(name=elem.name))
return params
return []
def _build_extracted_function(
key: str,
fn: JsFunctionExpression,
payload_id: str,
) -> JsFunctionDeclaration | None:
"""
Convert a dispatcher function-table entry into a standalone
`refinery.lib.scripts.js.model.JsFunctionDeclaration` of the same kind. Extracts parameters from
the payload destructuring and removes that statement.
The kind is carried rather than defaulted: an entry that was `async` returns a promise and one
that was a generator returns an iterator, and a declaration that dropped either would compute a
different value — a `yield` left in a body no longer marked `*` does not even parse.
The body is cloned before anything is taken out of it, so that building a declaration leaves
the entry it was built from exactly as it stood. Reading the statements out of the entry itself
would make extraction destructive, and extraction is attempted for every entry of a table
before any of them is installed: a later entry this pass cannot read would then abandon the
unwrap over a table whose earlier entries had already had their payload destructuring taken
away, which is a dispatcher whose callees name parameters nothing declares.
"""
params = _extract_params(fn, payload_id)
if params is None:
return None
if not isinstance(fn.body, JsBlockStatement):
return None
body = _clone_node(fn.body)
new_body_stmts = list(body.body)
if new_body_stmts and params:
first = new_body_stmts[0]
if isinstance(first, JsVariableDeclaration):
remaining = [
d for d in first.declarations
if not (
isinstance(d, JsVariableDeclarator)
and isinstance(d.id, JsArrayPattern)
and isinstance(d.init, JsIdentifier)
and d.init.name == payload_id
)
]
if not remaining:
new_body_stmts = new_body_stmts[1:]
else:
first.declarations = remaining
new_body = JsBlockStatement(body=new_body_stmts)
decl = JsFunctionDeclaration(
id=JsIdentifier(name=key),
params=list(params),
body=new_body,
generator=is_generator_function(fn),
is_async=is_async_function(fn),
)
return decl
def _is_object_create_null(node: Node) -> bool:
"""
Check if *node* is `Object.create(null)`.
"""
if not isinstance(node, JsCallExpression):
return False
if len(node.arguments) != 1 or not isinstance(node.arguments[0], JsNullLiteral):
return False
callee = node.callee
if not isinstance(callee, JsMemberExpression):
return False
if not isinstance(callee.object, JsIdentifier) or callee.object.name != 'Object':
return False
prop = callee.property
if isinstance(prop, JsStringLiteral):
return prop.value == 'create'
if isinstance(prop, JsIdentifier) and not callee.computed:
return prop.name == 'create'
return False
class JsDispatcherUnwrapper(ScopeProcessingTransformer):
"""
Detect and unwrap a dispatcher pattern. For each dispatcher found, extract the wrapped
functions, rewrite call sites, and remove the dispatcher scaffolding.
"""
def __init__(self):
super().__init__()
self._root: JsScript | None = None
def visit_JsScript(self, node: JsScript):
self._root = node
return super().visit_JsScript(node)
def _process_scope_body(self, scope: Node, body: list) -> None:
for func in list(body):
if not isinstance(func, JsFunctionDeclaration):
continue
info = _detect_dispatcher(func)
if info is None:
continue
self._unwrap_dispatcher(scope, body, info)
def _unwrap_dispatcher(
self,
scope: Node,
body: list,
info: _DispatcherInfo,
) -> None:
plan = self._plan_call_sites(scope, info)
if not self._plan_covers_every_reference(info, plan):
return
extracted: dict[str, JsFunctionDeclaration] = {}
for key, fn in info.fns_map.items():
decl = _build_extracted_function(key, fn, info.payload_id)
if decl is None:
return
extracted[key] = decl
for site in plan:
_replace_in_parent(site.node, site.replacement())
insert_idx = body.index(info.decl)
body.remove(info.decl)
for i, (key, decl) in enumerate(extracted.items()):
decl.parent = scope
body.insert(insert_idx + i, decl)
self.mark_changed()
self._remove_boilerplate(scope, body, info)
def _plan_call_sites(
self,
scope: Node,
info: _DispatcherInfo,
) -> list[_DispatchSite]:
"""
Every dispatch through *info* that this pass can read, with nothing replaced yet. Whether
the dispatcher may be removed at all is a question about the whole set, so the set has to
exist before the first replacement does.
A site nested inside another is kept rather than dropped, because the outer replacement does
not always take the inner one with it: the direct-call form carries the payload elements
into the call it builds, so a dispatch written into a payload survives the rewrite of the
dispatch it is an argument of. The order the walk yields is what makes both replaceable —
an ancestor comes first, and building its replacement adopts the arguments it reuses, so the
inner node is still reachable from its new holder when its own turn comes.
"""
planned: list[_DispatchSite] = []
for node in list(scope.walk()):
if isinstance(node, JsSequenceExpression):
site = self._read_direct_call(node, info)
elif isinstance(node, JsMemberExpression):
site = self._read_wrapped_ref(node, info)
elif isinstance(node, JsCallExpression):
site = self._read_bare_call(node, info)
else:
continue
if site is not None:
planned.append(site)
return planned
def _plan_covers_every_reference(
self,
info: _DispatcherInfo,
plan: list[_DispatchSite],
) -> bool:
"""
Whether *plan* replaces every reference to the dispatcher, so that removing its declaration
leaves nothing naming it. A dispatch this pass cannot read is one it is right to leave
alone, and what it leaves alone still calls the function it is leaving, so removing the
declaration anyway hands back a file that throws a `ReferenceError` where the original ran.
The whole unwrap is refused rather than the removal alone, because a dispatcher left
standing beside the extracted functions would still have to route through its own table,
and the payload the surviving dispatch writes is read by no extracted body.
What each site accounts for is the one occurrence of the name it consumes, not everything
standing inside it. A dispatch this pass cannot read is often written *within* one it can,
a payload argument being the ordinary place for a call, and the replacement carries such an
argument over rather than discarding it, so counting a subtree as covered would clear the
very reference that survives.
The dispatcher's own body is excluded rather than counted, since it goes with the
declaration. Everything else is asked of the model, so a same-named binding in another
scope is not mistaken for a use of this one.
"""
assert self._root is not None
if not isinstance(info.decl.id, JsIdentifier):
return False
model = model_cache(self, self._root).model
binding = model.binding_of(info.decl.id)
return not binding_has_references(
model,
binding,
exclude=info.decl,
exclude_ids={id(site.reference) for site in plan},
)
def _read_direct_call(
self,
seq: JsSequenceExpression,
info: _DispatcherInfo,
) -> _DispatchSite | None:
"""
The sequence expression dispatch call *seq* is:
(payload = [args], dispatcher("key")) -> key(args)
Also reads the wrapped variant where the return value is unwrapped via a member access
on the wrap key:
(payload = [args], dispatcher("key", s, wrapFlag)["wk"])
A `new` dispatch is read only in that wrapped variant. `new` hands back the object the
dispatcher returned only where it returned one, and the wrapper is the one branch that
does: everywhere else `new` yields the fresh instance and the result the call computed is
thrown away, which the call this would build hands back instead.
"""
if len(seq.expressions) != 2:
return None
assign, second = seq.expressions
if not isinstance(assign, JsAssignmentExpression):
return None
if assign.operator != '=':
return None
if not isinstance(assign.left, JsIdentifier) or assign.left.name != info.payload_id:
return None
if not isinstance(assign.right, JsArrayExpression):
return None
read = self._unwrap_dispatch_call(second, info)
if read is None:
return None
dispatch_call, through_the_wrap_key = read
if isinstance(dispatch_call, JsNewExpression) and not through_the_wrap_key:
return None
if not dispatch_call.arguments:
return None
key_arg = dispatch_call.arguments[0]
if not isinstance(key_arg, JsStringLiteral):
return None
if key_arg.value not in info.fns_map:
return None
if not self._flags_agree_with_the_reading(
dispatch_call,
info,
unwrapped=through_the_wrap_key,
selects_without_calling=False,
carries_the_payload=True,
):
return None
elements = assign.right.elements
if any(element is None for element in elements) and not self._a_hole_reads_undefined():
return None
args = [
make_undefined_expression() if e is None else e
for e in elements
]
assert isinstance(dispatch_call.callee, JsIdentifier)
return _DispatchSite(seq, dispatch_call.callee, key_arg.value, args)
def _a_hole_reads_undefined(self) -> bool:
"""
Whether a payload position written with no element in it reads `undefined`, which is what
spelling it out as `undefined` at the call site claims. A hole is not an element whose value
is `undefined`: the callee reaches it by reading that index off the payload array, so what
it finds is whatever `Array.prototype` answers there, and a file that wrote that prototype
answers something else.
"""
assert self._root is not None
return model_cache(self, self._root).effects.chain_roots_unwritten(list)
def _flags_agree_with_the_reading(
self,
call: JsCallExpression | JsNewExpression,
info: _DispatcherInfo,
*,
unwrapped: bool,
selects_without_calling: bool,
carries_the_payload: bool,
) -> bool:
"""
Whether the arguments after the key say the dispatch is the one the site reading it built.
The key alone selects the table entry; which of the dispatcher's branches runs, and what the
expression standing at the site therefore denotes, is decided by the two flags behind it.
Three readings can disagree with the flags, and each is a value the replacement would get
wrong rather than a shape it cannot spell. A dispatch spelling the wrap flag denotes the
wrapper object, so the access on the wrap key belongs to it and a reading without one hands
back the result the wrapper held. A dispatch spelling the create flag is handed the table
entry rather than what calling it returned, which is the wrapped-reference form's whole
premise and the ruin of every other one. And a dispatch spelling the init flag reaches its
callee with the payload emptied, so a reading that carries the payload elements into a call
passes arguments the original threw away.
Every argument beyond the third is dropped by the replacement, so it has to be one nothing
can miss. An argument this pass cannot read as a fixed string is refused outright in the
two flag positions, since such an argument may be any flag and may run something on the way
to being one.
"""
assert self._root is not None
flag = _flag_argument(call, 1)
rtype = _flag_argument(call, 2)
if isinstance(flag, _Unreadable) or isinstance(rtype, _Unreadable):
return False
if (info.wrap_flag is not None and rtype == info.wrap_flag) is not unwrapped:
return False
selects = info.create_flag is not None and flag == info.create_flag
if selects is not selects_without_calling:
return False
if carries_the_payload and flag == info.init_flag:
return False
effects = model_cache(self, self._root).effects
return all(
effects.is_side_effect_free(argument, discarded=True)
for argument in call.arguments[3:]
)
@staticmethod
def _unwrap_dispatch_call(
node: Node,
info: _DispatcherInfo,
) -> tuple[JsCallExpression | JsNewExpression, bool] | None:
"""
Extract a dispatcher call from *node*, which may be a bare call or a member access of
the form:
dispatcher(...)["wrapKey"]
Returns the call node together with whether the wrap-key access was read off it, or `None`.
The second half is what the caller checks the dispatch's own return-type flag against: the
access and the flag asking for the object it reads from are one fact written twice, and a
site holding one without the other denotes something else entirely.
"""
call = strip_parens(node)
through_the_wrap_key = False
if isinstance(call, JsMemberExpression) and info.wrap_key is not None:
if access_key(call) == info.wrap_key:
call = strip_parens(call.object)
through_the_wrap_key = True
if not isinstance(call, (JsCallExpression, JsNewExpression)):
return None
if not isinstance(call.callee, JsIdentifier):
return None
if call.callee.name != info.dispatcher_id:
return None
return call, through_the_wrap_key
def _read_wrapped_ref(
self,
member: JsMemberExpression,
info: _DispatcherInfo,
) -> _DispatchSite | None:
"""
The new-expression dispatch with wrap key access *member* is, which names the function it
selects rather than calling it:
new dispatcher("key", createFlag, wrapFlag)["wrapKey"] -> key
Both flags are what make it that. The create flag is what has the dispatcher hand the table
entry back instead of calling it, so without it this member denotes a *result*, and the wrap
flag is what puts that entry under the key being read.
"""
if info.wrap_key is None:
return None
if access_key(member) != info.wrap_key:
return None
new_expr = strip_parens(member.object)
if not isinstance(new_expr, JsNewExpression):
return None
if not isinstance(new_expr.callee, JsIdentifier):
return None
if new_expr.callee.name != info.dispatcher_id:
return None
if not new_expr.arguments:
return None
key_arg = new_expr.arguments[0]
if not isinstance(key_arg, JsStringLiteral):
return None
if key_arg.value not in info.fns_map:
return None
if not self._flags_agree_with_the_reading(
new_expr,
info,
unwrapped=True,
selects_without_calling=True,
carries_the_payload=False,
):
return None
return _DispatchSite(member, new_expr.callee, key_arg.value, None)
def _read_bare_call(
self,
call: JsCallExpression,
info: _DispatcherInfo,
) -> _DispatchSite | None:
"""
The bare `dispatcher("key")` call *call* is. These occur without a preceding payload
assignment, when the dispatched function takes no arguments.
A call standing second in a sequence expression is not one of them however it reads here:
the assignment in front of it is what fills the payload its callee takes its arguments from,
so rewriting it alone would call that callee with none. It belongs to `_read_direct_call`,
which reads the pair, and is left for that one to plan or to refuse whole. The parent is
read through any parentheses around the call, since a file that writes the grouping means
the same dispatch by it.
A call whose result is unwrapped on the wrap key is not one of them either: what that
member expression denotes is the callee's return value, so replacing the call alone leaves
the unwrap standing over a value that carries no such key. This pass has no reading of that
form, and an unread dispatch is what `_plan_covers_every_reference` refuses the unwrap on.
The zero arguments this reading gives its callee are a claim about the payload, not about
the site: the callee reads its arguments off the payload array, which a dispatch that does
not spell the init flag leaves holding whatever the last one put there. So an entry that
reads the payload at all is planned here only behind that flag, and an entry that reads
none is planned whatever the payload holds.
"""
parent = enclosing_operator(call)
if isinstance(parent, JsSequenceExpression):
return None
if (
info.wrap_key is not None
and isinstance(parent, JsMemberExpression)
and strip_parens(parent.object) is call
and access_key(parent) == info.wrap_key
):
return None
if not isinstance(call.callee, JsIdentifier):
return None
if call.callee.name != info.dispatcher_id:
return None
if not call.arguments:
return None
key_arg = call.arguments[0]
if not isinstance(key_arg, JsStringLiteral):
return None
if key_arg.value not in info.fns_map:
return None
if not self._flags_agree_with_the_reading(
call,
info,
unwrapped=False,
selects_without_calling=False,
carries_the_payload=False,
):
return None
if (
_flag_argument(call, 1) != info.init_flag
and _reads_the_payload(info.fns_map[key_arg.value], info.payload_id)
):
return None
return _DispatchSite(call, call.callee, key_arg.value, [])
def _remove_boilerplate(self, scope: Node, body: list, info: _DispatcherInfo) -> None:
"""
Remove dispatcher-related boilerplate declarations from the scope body.
"""
assert self._root is not None
model = model_cache(self, self._root).model
to_remove = []
for stmt in list(body):
if isinstance(stmt, JsVariableDeclaration):
for decl in stmt.declarations:
if not isinstance(decl, JsVariableDeclarator):
continue
if not isinstance(decl.id, JsIdentifier):
continue
if decl.id.name == info.payload_id and decl.init is None:
remove_declarator(decl)
break
if info.cache_id and decl.id.name == info.cache_id:
if decl.init is not None and _is_object_create_null(decl.init):
remove_declarator(decl)
break
elif isinstance(stmt, JsFunctionDeclaration):
if (
isinstance(stmt.id, JsIdentifier)
and isinstance(stmt.body, JsBlockStatement)
and not stmt.body.body
and not stmt.params
):
binding = model.binding_of(stmt.id)
if not binding_has_references(model, binding, exclude=stmt):
to_remove.append(stmt)
for stmt in to_remove:
_remove_from_parent(stmt)
Classes
class JsDispatcherUnwrapper-
Detect and unwrap a dispatcher pattern. For each dispatcher found, extract the wrapped functions, rewrite call sites, and remove the dispatcher scaffolding.
Expand source code Browse git
class JsDispatcherUnwrapper(ScopeProcessingTransformer): """ Detect and unwrap a dispatcher pattern. For each dispatcher found, extract the wrapped functions, rewrite call sites, and remove the dispatcher scaffolding. """ def __init__(self): super().__init__() self._root: JsScript | None = None def visit_JsScript(self, node: JsScript): self._root = node return super().visit_JsScript(node) def _process_scope_body(self, scope: Node, body: list) -> None: for func in list(body): if not isinstance(func, JsFunctionDeclaration): continue info = _detect_dispatcher(func) if info is None: continue self._unwrap_dispatcher(scope, body, info) def _unwrap_dispatcher( self, scope: Node, body: list, info: _DispatcherInfo, ) -> None: plan = self._plan_call_sites(scope, info) if not self._plan_covers_every_reference(info, plan): return extracted: dict[str, JsFunctionDeclaration] = {} for key, fn in info.fns_map.items(): decl = _build_extracted_function(key, fn, info.payload_id) if decl is None: return extracted[key] = decl for site in plan: _replace_in_parent(site.node, site.replacement()) insert_idx = body.index(info.decl) body.remove(info.decl) for i, (key, decl) in enumerate(extracted.items()): decl.parent = scope body.insert(insert_idx + i, decl) self.mark_changed() self._remove_boilerplate(scope, body, info) def _plan_call_sites( self, scope: Node, info: _DispatcherInfo, ) -> list[_DispatchSite]: """ Every dispatch through *info* that this pass can read, with nothing replaced yet. Whether the dispatcher may be removed at all is a question about the whole set, so the set has to exist before the first replacement does. A site nested inside another is kept rather than dropped, because the outer replacement does not always take the inner one with it: the direct-call form carries the payload elements into the call it builds, so a dispatch written into a payload survives the rewrite of the dispatch it is an argument of. The order the walk yields is what makes both replaceable — an ancestor comes first, and building its replacement adopts the arguments it reuses, so the inner node is still reachable from its new holder when its own turn comes. """ planned: list[_DispatchSite] = [] for node in list(scope.walk()): if isinstance(node, JsSequenceExpression): site = self._read_direct_call(node, info) elif isinstance(node, JsMemberExpression): site = self._read_wrapped_ref(node, info) elif isinstance(node, JsCallExpression): site = self._read_bare_call(node, info) else: continue if site is not None: planned.append(site) return planned def _plan_covers_every_reference( self, info: _DispatcherInfo, plan: list[_DispatchSite], ) -> bool: """ Whether *plan* replaces every reference to the dispatcher, so that removing its declaration leaves nothing naming it. A dispatch this pass cannot read is one it is right to leave alone, and what it leaves alone still calls the function it is leaving, so removing the declaration anyway hands back a file that throws a `ReferenceError` where the original ran. The whole unwrap is refused rather than the removal alone, because a dispatcher left standing beside the extracted functions would still have to route through its own table, and the payload the surviving dispatch writes is read by no extracted body. What each site accounts for is the one occurrence of the name it consumes, not everything standing inside it. A dispatch this pass cannot read is often written *within* one it can, a payload argument being the ordinary place for a call, and the replacement carries such an argument over rather than discarding it, so counting a subtree as covered would clear the very reference that survives. The dispatcher's own body is excluded rather than counted, since it goes with the declaration. Everything else is asked of the model, so a same-named binding in another scope is not mistaken for a use of this one. """ assert self._root is not None if not isinstance(info.decl.id, JsIdentifier): return False model = model_cache(self, self._root).model binding = model.binding_of(info.decl.id) return not binding_has_references( model, binding, exclude=info.decl, exclude_ids={id(site.reference) for site in plan}, ) def _read_direct_call( self, seq: JsSequenceExpression, info: _DispatcherInfo, ) -> _DispatchSite | None: """ The sequence expression dispatch call *seq* is: (payload = [args], dispatcher("key")) -> key(args) Also reads the wrapped variant where the return value is unwrapped via a member access on the wrap key: (payload = [args], dispatcher("key", s, wrapFlag)["wk"]) A `new` dispatch is read only in that wrapped variant. `new` hands back the object the dispatcher returned only where it returned one, and the wrapper is the one branch that does: everywhere else `new` yields the fresh instance and the result the call computed is thrown away, which the call this would build hands back instead. """ if len(seq.expressions) != 2: return None assign, second = seq.expressions if not isinstance(assign, JsAssignmentExpression): return None if assign.operator != '=': return None if not isinstance(assign.left, JsIdentifier) or assign.left.name != info.payload_id: return None if not isinstance(assign.right, JsArrayExpression): return None read = self._unwrap_dispatch_call(second, info) if read is None: return None dispatch_call, through_the_wrap_key = read if isinstance(dispatch_call, JsNewExpression) and not through_the_wrap_key: return None if not dispatch_call.arguments: return None key_arg = dispatch_call.arguments[0] if not isinstance(key_arg, JsStringLiteral): return None if key_arg.value not in info.fns_map: return None if not self._flags_agree_with_the_reading( dispatch_call, info, unwrapped=through_the_wrap_key, selects_without_calling=False, carries_the_payload=True, ): return None elements = assign.right.elements if any(element is None for element in elements) and not self._a_hole_reads_undefined(): return None args = [ make_undefined_expression() if e is None else e for e in elements ] assert isinstance(dispatch_call.callee, JsIdentifier) return _DispatchSite(seq, dispatch_call.callee, key_arg.value, args) def _a_hole_reads_undefined(self) -> bool: """ Whether a payload position written with no element in it reads `undefined`, which is what spelling it out as `undefined` at the call site claims. A hole is not an element whose value is `undefined`: the callee reaches it by reading that index off the payload array, so what it finds is whatever `Array.prototype` answers there, and a file that wrote that prototype answers something else. """ assert self._root is not None return model_cache(self, self._root).effects.chain_roots_unwritten(list) def _flags_agree_with_the_reading( self, call: JsCallExpression | JsNewExpression, info: _DispatcherInfo, *, unwrapped: bool, selects_without_calling: bool, carries_the_payload: bool, ) -> bool: """ Whether the arguments after the key say the dispatch is the one the site reading it built. The key alone selects the table entry; which of the dispatcher's branches runs, and what the expression standing at the site therefore denotes, is decided by the two flags behind it. Three readings can disagree with the flags, and each is a value the replacement would get wrong rather than a shape it cannot spell. A dispatch spelling the wrap flag denotes the wrapper object, so the access on the wrap key belongs to it and a reading without one hands back the result the wrapper held. A dispatch spelling the create flag is handed the table entry rather than what calling it returned, which is the wrapped-reference form's whole premise and the ruin of every other one. And a dispatch spelling the init flag reaches its callee with the payload emptied, so a reading that carries the payload elements into a call passes arguments the original threw away. Every argument beyond the third is dropped by the replacement, so it has to be one nothing can miss. An argument this pass cannot read as a fixed string is refused outright in the two flag positions, since such an argument may be any flag and may run something on the way to being one. """ assert self._root is not None flag = _flag_argument(call, 1) rtype = _flag_argument(call, 2) if isinstance(flag, _Unreadable) or isinstance(rtype, _Unreadable): return False if (info.wrap_flag is not None and rtype == info.wrap_flag) is not unwrapped: return False selects = info.create_flag is not None and flag == info.create_flag if selects is not selects_without_calling: return False if carries_the_payload and flag == info.init_flag: return False effects = model_cache(self, self._root).effects return all( effects.is_side_effect_free(argument, discarded=True) for argument in call.arguments[3:] ) @staticmethod def _unwrap_dispatch_call( node: Node, info: _DispatcherInfo, ) -> tuple[JsCallExpression | JsNewExpression, bool] | None: """ Extract a dispatcher call from *node*, which may be a bare call or a member access of the form: dispatcher(...)["wrapKey"] Returns the call node together with whether the wrap-key access was read off it, or `None`. The second half is what the caller checks the dispatch's own return-type flag against: the access and the flag asking for the object it reads from are one fact written twice, and a site holding one without the other denotes something else entirely. """ call = strip_parens(node) through_the_wrap_key = False if isinstance(call, JsMemberExpression) and info.wrap_key is not None: if access_key(call) == info.wrap_key: call = strip_parens(call.object) through_the_wrap_key = True if not isinstance(call, (JsCallExpression, JsNewExpression)): return None if not isinstance(call.callee, JsIdentifier): return None if call.callee.name != info.dispatcher_id: return None return call, through_the_wrap_key def _read_wrapped_ref( self, member: JsMemberExpression, info: _DispatcherInfo, ) -> _DispatchSite | None: """ The new-expression dispatch with wrap key access *member* is, which names the function it selects rather than calling it: new dispatcher("key", createFlag, wrapFlag)["wrapKey"] -> key Both flags are what make it that. The create flag is what has the dispatcher hand the table entry back instead of calling it, so without it this member denotes a *result*, and the wrap flag is what puts that entry under the key being read. """ if info.wrap_key is None: return None if access_key(member) != info.wrap_key: return None new_expr = strip_parens(member.object) if not isinstance(new_expr, JsNewExpression): return None if not isinstance(new_expr.callee, JsIdentifier): return None if new_expr.callee.name != info.dispatcher_id: return None if not new_expr.arguments: return None key_arg = new_expr.arguments[0] if not isinstance(key_arg, JsStringLiteral): return None if key_arg.value not in info.fns_map: return None if not self._flags_agree_with_the_reading( new_expr, info, unwrapped=True, selects_without_calling=True, carries_the_payload=False, ): return None return _DispatchSite(member, new_expr.callee, key_arg.value, None) def _read_bare_call( self, call: JsCallExpression, info: _DispatcherInfo, ) -> _DispatchSite | None: """ The bare `dispatcher("key")` call *call* is. These occur without a preceding payload assignment, when the dispatched function takes no arguments. A call standing second in a sequence expression is not one of them however it reads here: the assignment in front of it is what fills the payload its callee takes its arguments from, so rewriting it alone would call that callee with none. It belongs to `_read_direct_call`, which reads the pair, and is left for that one to plan or to refuse whole. The parent is read through any parentheses around the call, since a file that writes the grouping means the same dispatch by it. A call whose result is unwrapped on the wrap key is not one of them either: what that member expression denotes is the callee's return value, so replacing the call alone leaves the unwrap standing over a value that carries no such key. This pass has no reading of that form, and an unread dispatch is what `_plan_covers_every_reference` refuses the unwrap on. The zero arguments this reading gives its callee are a claim about the payload, not about the site: the callee reads its arguments off the payload array, which a dispatch that does not spell the init flag leaves holding whatever the last one put there. So an entry that reads the payload at all is planned here only behind that flag, and an entry that reads none is planned whatever the payload holds. """ parent = enclosing_operator(call) if isinstance(parent, JsSequenceExpression): return None if ( info.wrap_key is not None and isinstance(parent, JsMemberExpression) and strip_parens(parent.object) is call and access_key(parent) == info.wrap_key ): return None if not isinstance(call.callee, JsIdentifier): return None if call.callee.name != info.dispatcher_id: return None if not call.arguments: return None key_arg = call.arguments[0] if not isinstance(key_arg, JsStringLiteral): return None if key_arg.value not in info.fns_map: return None if not self._flags_agree_with_the_reading( call, info, unwrapped=False, selects_without_calling=False, carries_the_payload=False, ): return None if ( _flag_argument(call, 1) != info.init_flag and _reads_the_payload(info.fns_map[key_arg.value], info.payload_id) ): return None return _DispatchSite(call, call.callee, key_arg.value, []) def _remove_boilerplate(self, scope: Node, body: list, info: _DispatcherInfo) -> None: """ Remove dispatcher-related boilerplate declarations from the scope body. """ assert self._root is not None model = model_cache(self, self._root).model to_remove = [] for stmt in list(body): if isinstance(stmt, JsVariableDeclaration): for decl in stmt.declarations: if not isinstance(decl, JsVariableDeclarator): continue if not isinstance(decl.id, JsIdentifier): continue if decl.id.name == info.payload_id and decl.init is None: remove_declarator(decl) break if info.cache_id and decl.id.name == info.cache_id: if decl.init is not None and _is_object_create_null(decl.init): remove_declarator(decl) break elif isinstance(stmt, JsFunctionDeclaration): if ( isinstance(stmt.id, JsIdentifier) and isinstance(stmt.body, JsBlockStatement) and not stmt.body.body and not stmt.params ): binding = model.binding_of(stmt.id) if not binding_has_references(model, binding, exclude=stmt): to_remove.append(stmt) for stmt in to_remove: _remove_from_parent(stmt)Ancestors
Methods
def visit_JsScript(self, node)-
Expand source code Browse git
def visit_JsScript(self, node: JsScript): self._root = node return super().visit_JsScript(node)
Inherited members