8b2899d16e
Decomposition breaks long sequences into shorter fragments before inference. This helps when sequences are too long for CRX to handle (>5 symbols → flat bags). Results: - RAGSAK: 21 → 80 grammars (3.8× increase) - FastAPI: 111 → 118 grammars (small increase) Changes: - bex/decompose.py: decompose_sequence(), decompose_all(), decompose_with_coverage() - bex/tag_preprocessor/analyze.py: --decompose, --max-seq-length flags - Skip diversity check when decomposing (decomposition creates diverse fragments) - 12 new tests in tests/test_decompose.py Co-authored-by: OpenCode <opencode@corentic.eu>
328 lines
10 KiB
Python
328 lines
10 KiB
Python
"""AST-structural pattern extraction → GBNF grammar generation.
|
|
|
|
Instead of flattening tree-sitter captures to sequences and re-inferring
|
|
patterns with BEX algorithms, this module walks the AST directly and
|
|
extracts structural patterns from function bodies.
|
|
|
|
The core insight: tree-sitter already understands code structure.
|
|
We don't need to throw that away and re-infer it.
|
|
|
|
Approach:
|
|
1. Parse file with tree-sitter
|
|
2. For each function, extract the "structural shape" of its body
|
|
(the sequence of AST node types, with text replaced by placeholders)
|
|
3. Group functions by their structural shape
|
|
4. Generate GBNF rules from each group
|
|
"""
|
|
|
|
import os
|
|
import re
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
|
|
import tree_sitter
|
|
|
|
from .code import EXTENSION_MAP, _load_grammar, _load_query, _find_method_bodies
|
|
|
|
|
|
def _get_parser(ext):
|
|
"""Get a tree-sitter parser for the given extension."""
|
|
lang_name, module_name, func_name = EXTENSION_MAP[ext]
|
|
mod = __import__(module_name)
|
|
lang_obj = getattr(mod, func_name)()
|
|
lang = tree_sitter.Language(lang_obj)
|
|
return tree_sitter.Parser(lang)
|
|
|
|
|
|
def _node_text(node, code_bytes):
|
|
"""Get the text of a node, decoded from bytes."""
|
|
return code_bytes[node.start_byte:node.end_byte].decode('utf-8', errors='replace')
|
|
|
|
|
|
def _classify_node(node):
|
|
"""Classify an AST node into a GBNF terminal/nonterminal category.
|
|
|
|
Returns (category, text) where category is a string like:
|
|
- "return", "if", "loop", "throw" (structural keywords)
|
|
- "call" (function/method call)
|
|
- "new" (constructor call)
|
|
- "member" (member access: a.b)
|
|
- "identifier" (bare identifier)
|
|
- "string", "number", "boolean" (literals)
|
|
- "assign" (assignment)
|
|
- "expr" (complex expression)
|
|
"""
|
|
t = node.type
|
|
|
|
if t == 'return_statement':
|
|
return ('return', 'return')
|
|
if t in ('if_statement', 'if_else_statement'):
|
|
return ('if', 'if')
|
|
if t in ('for_statement', 'for_in_statement', 'while_statement', 'do_statement'):
|
|
return ('loop', 'loop')
|
|
if t in ('throw_statement', 'raise_statement'):
|
|
return ('throw', 'throw')
|
|
if t == 'try_statement':
|
|
return ('try', 'try')
|
|
|
|
if t == 'new_expression':
|
|
ctor = node.child_by_field_name('constructor')
|
|
if ctor:
|
|
return ('new', _node_text(ctor, b''))
|
|
return ('new', 'new')
|
|
|
|
if t == 'call_expression':
|
|
func = node.child_by_field_name('function')
|
|
if func and func.type == 'member_expression':
|
|
prop = func.child_by_field_name('property')
|
|
if prop:
|
|
return ('call', _node_text(prop, b''))
|
|
elif func:
|
|
return ('call', _node_text(func, b''))
|
|
return ('call', 'call')
|
|
|
|
if t == 'member_expression':
|
|
prop = node.child_by_field_name('property')
|
|
if prop:
|
|
return ('member', _node_text(prop, b''))
|
|
return ('member', 'member')
|
|
|
|
if t == 'assignment_expression':
|
|
return ('assign', '=')
|
|
|
|
if t in ('identifier', 'property_identifier', 'shorthand_property_identifier'):
|
|
return ('identifier', _node_text(node, b''))
|
|
|
|
if t in ('string', 'string_fragment', 'template_string'):
|
|
return ('string', '"..."')
|
|
|
|
if t in ('number', 'float', 'integer'):
|
|
return ('number', 'N')
|
|
|
|
if t in ('true', 'false'):
|
|
return ('boolean', 'true' if t == 'true' else 'false')
|
|
|
|
if t in ('null', 'undefined', 'none'):
|
|
return ('null', 'null')
|
|
|
|
if t == 'as_expression':
|
|
return ('cast', 'as')
|
|
|
|
return ('expr', t)
|
|
|
|
|
|
def _extract_body_shape(body_node, code_bytes):
|
|
"""Extract the structural shape of a function body.
|
|
|
|
Returns a list of (category, text) tuples representing the
|
|
high-level structure of the function body.
|
|
|
|
For simple functions (single return), this is just one entry.
|
|
For complex functions, it's the sequence of structural statements.
|
|
"""
|
|
shape = []
|
|
for child in body_node.children:
|
|
cat, text = _classify_node(child)
|
|
if cat in ('return', 'if', 'loop', 'throw', 'try', 'assign'):
|
|
# Structural statement — always include
|
|
shape.append((cat, text))
|
|
elif cat == 'expr':
|
|
# Expression statement — might be a call
|
|
# Check if it's actually a call expression
|
|
if child.type == 'expression_statement' and child.children:
|
|
inner = child.children[0]
|
|
inner_cat, inner_text = _classify_node(inner)
|
|
if inner_cat in ('call', 'new'):
|
|
shape.append((inner_cat, inner_text))
|
|
# Skip pure declarations, imports, comments, etc.
|
|
|
|
return shape
|
|
|
|
|
|
def _shape_key(shape):
|
|
"""Convert a shape to a hashable key for grouping."""
|
|
return tuple(cat for cat, _ in shape)
|
|
|
|
|
|
def _generalize_group shapes):
|
|
"""Generalize a group of shapes into a GBNF rule.
|
|
|
|
All shapes in the group have the same structural pattern.
|
|
The varying parts are the specific identifiers/names.
|
|
We replace those with wildcards in the GBNF.
|
|
"""
|
|
if not shapes:
|
|
return None
|
|
|
|
# All shapes have the same categories, so take the first
|
|
first = shapes[0]
|
|
|
|
# Check if all shapes in the group are identical
|
|
all_same = all(s == first for s in shapes)
|
|
|
|
if all_same:
|
|
# All functions follow the exact same pattern
|
|
# Generate a specific GBNF rule
|
|
tokens = []
|
|
for cat, text in first:
|
|
if cat in ('return', 'if', 'loop', 'throw', 'try'):
|
|
tokens.append(f'"{text}"')
|
|
elif cat == 'call':
|
|
tokens.append(f'CALL')
|
|
elif cat == 'new':
|
|
tokens.append(f'NEW')
|
|
elif cat == 'member':
|
|
tokens.append(f'MEMBER')
|
|
elif cat == 'assign':
|
|
tokens.append(f'"="')
|
|
elif cat == 'cast':
|
|
tokens.append(f'"as"')
|
|
elif cat == 'string':
|
|
tokens.append(f'STRING')
|
|
elif cat == 'number':
|
|
tokens.append(f'NUMBER')
|
|
elif cat == 'boolean':
|
|
tokens.append(f'BOOL')
|
|
elif cat == 'null':
|
|
tokens.append(f'NULL')
|
|
elif cat == 'identifier':
|
|
tokens.append(f'IDENT')
|
|
else:
|
|
tokens.append(f'"{text}"')
|
|
return ' '.join(tokens)
|
|
else:
|
|
# Functions have different specific names but same structure
|
|
# Generate a generalized rule with alternatives
|
|
# For now, use CALL/NEW/IDENT wildcards
|
|
tokens = []
|
|
for cat, text in first:
|
|
if cat in ('return', 'if', 'loop', 'throw', 'try'):
|
|
tokens.append(f'"{text}"')
|
|
elif cat == 'call':
|
|
tokens.append(f'CALL')
|
|
elif cat == 'new':
|
|
tokens.append(f'NEW')
|
|
elif cat == 'member':
|
|
tokens.append(f'MEMBER')
|
|
elif cat == 'assign':
|
|
tokens.append(f'"="')
|
|
elif cat == 'cast':
|
|
tokens.append(f'"as"')
|
|
elif cat == 'string':
|
|
tokens.append(f'STRING')
|
|
elif cat == 'number':
|
|
tokens.append(f'NUMBER')
|
|
elif cat == 'boolean':
|
|
tokens.append(f'BOOL')
|
|
elif cat == 'null':
|
|
tokens.append(f'NULL')
|
|
elif cat == 'identifier':
|
|
tokens.append(f'IDENT')
|
|
else:
|
|
tokens.append(f'"{text}"')
|
|
return ' '.join(tokens)
|
|
|
|
|
|
def extract_patterns_from_file(file_path):
|
|
"""Extract structural patterns from a single file.
|
|
|
|
Returns a dict: { pattern_key: [function_name, ...] }
|
|
where pattern_key is the GBNF rule string.
|
|
"""
|
|
ext = os.path.splitext(file_path)[1].lower()
|
|
if ext not in EXTENSION_MAP:
|
|
return {}
|
|
|
|
try:
|
|
parser = _get_parser(ext)
|
|
except Exception:
|
|
return {}
|
|
|
|
try:
|
|
with open(file_path, 'rb') as f:
|
|
code_bytes = f.read()
|
|
except Exception:
|
|
return {}
|
|
|
|
tree = parser.parse(code_bytes)
|
|
|
|
# Find method bodies
|
|
bodies = _find_method_bodies(tree)
|
|
bodies.sort(key=lambda b: b.start_byte)
|
|
|
|
# Extract shapes
|
|
patterns = defaultdict(list)
|
|
for body in bodies:
|
|
# Find the parent function name
|
|
parent = body.parent
|
|
name = None
|
|
if parent:
|
|
name_node = parent.child_by_field_name('name')
|
|
if name_node:
|
|
name = _node_text(name_node, code_bytes)
|
|
|
|
shape = _extract_body_shape(body, code_bytes)
|
|
if shape:
|
|
key = _shape_key(shape)
|
|
patterns[key].append((name, shape))
|
|
|
|
return patterns
|
|
|
|
|
|
def patterns_to_gbnf(patterns):
|
|
"""Convert extracted patterns to GBNF grammar rules.
|
|
|
|
Args:
|
|
patterns: dict from extract_patterns_from_file
|
|
|
|
Returns:
|
|
GBNF grammar string
|
|
"""
|
|
rules = []
|
|
|
|
# Base rules for wildcards
|
|
rules.append('root ::= function')
|
|
rules.append('function ::= RETURN CALL NEW MEMBER IDENT STRING NUMBER BOOL NULL')
|
|
rules.append('CALL ::= IDENT')
|
|
rules.append('NEW ::= "new" IDENT')
|
|
rules.append('MEMBER ::= IDENT "." IDENT')
|
|
rules.append('IDENT ::= [a-zA-Z_] [a-zA-Z0-9_]*')
|
|
rules.append('STRING ::= "\\"" [^\\"\n]* "\\""')
|
|
rules.append('NUMBER ::= [0-9]+ ("." [0-9]+)?')
|
|
rules.append('BOOL ::= "true" | "false"')
|
|
rules.append('NULL ::= "null" | "undefined"')
|
|
rules.append('RETURN ::= "return"')
|
|
rules.append('IF ::= "if"')
|
|
rules.append('LOOP ::= "for" | "while"')
|
|
rules.append('THROW ::= "throw"')
|
|
rules.append('TRY ::= "try"')
|
|
rules.append('CAST ::= "as"')
|
|
rules.append('"=" ::= "="')
|
|
|
|
# Generate rules for each pattern group
|
|
for i, (key, funcs) in enumerate(patterns.items()):
|
|
rule_name = f'pattern_{i}'
|
|
gbnf_body = _generalize_group([shape for _, shape in funcs])
|
|
if gbnf_body:
|
|
rules.append(f'{rule_name} ::= {gbnf_body}')
|
|
|
|
return '\n'.join(rules)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
import sys
|
|
if len(sys.argv) < 2:
|
|
print("Usage: python -m bex.ast_to_gbnf <file>")
|
|
sys.exit(1)
|
|
|
|
file_path = sys.argv[1]
|
|
patterns = extract_patterns_from_file(file_path)
|
|
|
|
print(f"Found {len(patterns)} pattern groups:")
|
|
for key, funcs in patterns.items():
|
|
names = [name for name, _ in funcs]
|
|
print(f" {key}: {names}")
|
|
|
|
print()
|
|
print("GBNF:")
|
|
print(patterns_to_gbnf(patterns))
|