310 lines
13 KiB
Python
310 lines
13 KiB
Python
"""Fortran extractor. Moved verbatim from graphify/extract.py."""
|
|
from __future__ import annotations
|
|
|
|
|
|
from pathlib import Path
|
|
from graphify.extractors.base import _file_stem, _make_id, _read_text
|
|
|
|
|
|
_FORTRAN_CPP_EXTS = {".F", ".F90", ".F95", ".F03", ".F08"}
|
|
|
|
def _cpp_preprocess(path: Path) -> bytes:
|
|
"""Run cpp -w -P on a capital-F Fortran file and return preprocessed bytes.
|
|
|
|
Falls back to raw file bytes if cpp is not available. Capital-F extensions
|
|
conventionally require C preprocessor expansion (#ifdef MPI, #define REAL8, etc.)
|
|
before parsing.
|
|
|
|
Security (F-007): we pass `-nostdinc` and `-I /dev/null` so a malicious
|
|
source file containing `#include "/home/victim/.ssh/id_rsa"` (or any other
|
|
include directive) cannot inline arbitrary host files into the output that
|
|
we then ship to an LLM. Without these flags `cpp` happily resolves any
|
|
relative or absolute include path it can read, which is a corpus-side
|
|
file-exfiltration vector.
|
|
"""
|
|
import shutil
|
|
import subprocess
|
|
if not shutil.which("cpp"):
|
|
return path.read_bytes()
|
|
try:
|
|
# Pass an absolute path so a corpus file named like "-I/etc/x.F90" cannot
|
|
# be parsed by cpp as an option (cpp does not accept a "--" end-of-options
|
|
# terminator). An absolute path always begins with "/".
|
|
result = subprocess.run(
|
|
["cpp", "-w", "-P", "-nostdinc", "-I", "/dev/null", str(path.resolve())],
|
|
capture_output=True,
|
|
timeout=30,
|
|
)
|
|
if result.returncode == 0 and result.stdout:
|
|
return result.stdout
|
|
except Exception:
|
|
pass
|
|
return path.read_bytes()
|
|
|
|
def extract_fortran(path: Path) -> dict:
|
|
"""Extract programs, modules, subroutines, functions, use statements, and calls from Fortran files.
|
|
|
|
Capital-F extensions (.F, .F90, etc.) are run through the C preprocessor before
|
|
parsing so #ifdef/#define macros are resolved.
|
|
"""
|
|
try:
|
|
import tree_sitter_fortran as tsfortran
|
|
from tree_sitter import Language, Parser
|
|
except ImportError:
|
|
return {"nodes": [], "edges": [], "error": "tree-sitter-fortran not installed"}
|
|
|
|
try:
|
|
language = Language(tsfortran.language())
|
|
parser = Parser(language)
|
|
source = _cpp_preprocess(path) if path.suffix in _FORTRAN_CPP_EXTS else path.read_bytes()
|
|
tree = parser.parse(source)
|
|
root = tree.root_node
|
|
except Exception as e:
|
|
return {"nodes": [], "edges": [], "error": str(e)}
|
|
|
|
stem = _file_stem(path)
|
|
str_path = str(path)
|
|
nodes: list[dict] = []
|
|
edges: list[dict] = []
|
|
seen_ids: set[str] = set()
|
|
scope_bodies: list[tuple[str, object]] = []
|
|
|
|
def add_node(nid: str, label: str, line: int) -> None:
|
|
if nid not in seen_ids:
|
|
seen_ids.add(nid)
|
|
nodes.append({
|
|
"id": nid,
|
|
"label": label,
|
|
"file_type": "code",
|
|
"source_file": str_path,
|
|
"source_location": f"L{line}",
|
|
})
|
|
|
|
def add_edge(src: str, tgt: str, relation: str, line: int,
|
|
confidence: str = "EXTRACTED", weight: float = 1.0,
|
|
context: str | None = None) -> None:
|
|
edge = {
|
|
"source": src,
|
|
"target": tgt,
|
|
"relation": relation,
|
|
"confidence": confidence,
|
|
"source_file": str_path,
|
|
"source_location": f"L{line}",
|
|
"weight": weight,
|
|
}
|
|
if context:
|
|
edge["context"] = context
|
|
edges.append(edge)
|
|
|
|
file_nid = _make_id(str(path))
|
|
add_node(file_nid, path.name, 1)
|
|
|
|
def _fortran_name(stmt_node) -> str | None:
|
|
"""Extract name from a *_statement node. Fortran is case-insensitive; lowercase."""
|
|
for child in stmt_node.children:
|
|
if child.type in ("name", "identifier"):
|
|
return _read_text(child, source).lower()
|
|
return None
|
|
|
|
def ensure_named_node(name: str, line: int) -> str:
|
|
nid = _make_id(stem, name)
|
|
if nid in seen_ids:
|
|
return nid
|
|
nid = _make_id(name)
|
|
if nid not in seen_ids:
|
|
# The name isn't defined in this file, so this is a cross-file reference
|
|
# (e.g. a `Thing` type annotation imported from another module). Emit a
|
|
# SOURCELESS stub — like the inheritance-base path below — so the
|
|
# corpus-level rewire can collapse it onto the real definition. A sourced
|
|
# stub here makes _disambiguate_colliding_node_ids bake the referencing
|
|
# file's path (with extension) into the id and blocks the rewire, which is
|
|
# the phantom-duplicate-node bug (#1402).
|
|
seen_ids.add(nid)
|
|
nodes.append({
|
|
"id": nid,
|
|
"label": name,
|
|
"file_type": "code",
|
|
"source_file": "",
|
|
"source_location": "",
|
|
"origin_file": str_path,
|
|
})
|
|
return nid
|
|
|
|
def emit_signature_refs(scope_node, fn_nid: str, is_function: bool) -> None:
|
|
"""Emit references[parameter_type] / references[return_type] edges for
|
|
a subroutine/function based on its variable_declaration siblings."""
|
|
stmt_type = "function_statement" if is_function else "subroutine_statement"
|
|
stmt = next((c for c in scope_node.children if c.type == stmt_type), None)
|
|
if stmt is None:
|
|
return
|
|
param_names: set[str] = set()
|
|
params_node = next((c for c in stmt.children if c.type == "parameters"), None)
|
|
if params_node is not None:
|
|
for c in params_node.children:
|
|
if c.type == "identifier":
|
|
param_names.add(_read_text(c, source).lower())
|
|
result_name: str | None = None
|
|
if is_function:
|
|
result_node = next((c for c in stmt.children if c.type == "function_result"), None)
|
|
if result_node is not None:
|
|
res_id = next((c for c in result_node.children if c.type == "identifier"), None)
|
|
if res_id is not None:
|
|
result_name = _read_text(res_id, source).lower()
|
|
else:
|
|
# implicit result variable: same name as the function
|
|
result_name = _fortran_name(stmt)
|
|
for child in scope_node.children:
|
|
if child.type != "variable_declaration":
|
|
continue
|
|
derived = next((c for c in child.children if c.type == "derived_type"), None)
|
|
if derived is None:
|
|
continue
|
|
type_name_node = next((c for c in derived.children if c.type == "type_name"), None)
|
|
if type_name_node is None:
|
|
continue
|
|
type_name = _read_text(type_name_node, source).lower()
|
|
for var in child.children:
|
|
if var.type != "identifier":
|
|
continue
|
|
var_name = _read_text(var, source).lower()
|
|
var_line = var.start_point[0] + 1
|
|
if var_name in param_names:
|
|
tgt = ensure_named_node(type_name, var_line)
|
|
if tgt != fn_nid:
|
|
add_edge(fn_nid, tgt, "references", var_line, context="parameter_type")
|
|
elif is_function and var_name == result_name:
|
|
tgt = ensure_named_node(type_name, var_line)
|
|
if tgt != fn_nid:
|
|
add_edge(fn_nid, tgt, "references", var_line, context="return_type")
|
|
|
|
def walk_calls(node, scope_nid: str) -> None:
|
|
if node is None:
|
|
return
|
|
t = node.type
|
|
if t in ("subroutine", "function", "module", "program", "internal_procedures"):
|
|
return
|
|
# call FOO(args) — tree-sitter-fortran uses subroutine_call
|
|
if t == "subroutine_call":
|
|
name_node = next((c for c in node.children if c.type == "identifier"), None)
|
|
if name_node:
|
|
callee = _read_text(name_node, source).lower()
|
|
target_nid = _make_id(stem, callee)
|
|
add_edge(scope_nid, target_nid, "calls", node.start_point[0] + 1,
|
|
confidence="EXTRACTED", context="call")
|
|
# x = compute(args) — function invocations are `call_expression`, which
|
|
# shares Fortran's `name(...)` syntax with array indexing. Only emit a
|
|
# call edge when the callee resolves to a procedure defined in this file
|
|
# (an array variable produces no matching node), so array accesses can't
|
|
# fabricate spurious `calls` edges.
|
|
elif t == "call_expression":
|
|
name_node = next((c for c in node.children if c.type == "identifier"), None)
|
|
if name_node:
|
|
callee = _read_text(name_node, source).lower()
|
|
target_nid = _make_id(stem, callee)
|
|
if target_nid in seen_ids and target_nid != scope_nid:
|
|
add_edge(scope_nid, target_nid, "calls", node.start_point[0] + 1,
|
|
confidence="EXTRACTED", context="call")
|
|
for child in node.children:
|
|
walk_calls(child, scope_nid)
|
|
|
|
def walk(node, scope_nid: str) -> None:
|
|
t = node.type
|
|
|
|
if t == "program":
|
|
stmt = next((c for c in node.children if c.type == "program_statement"), None)
|
|
name = _fortran_name(stmt) if stmt else None
|
|
if name:
|
|
nid = _make_id(stem, name)
|
|
line = node.start_point[0] + 1
|
|
add_node(nid, name, line)
|
|
add_edge(file_nid, nid, "defines", line)
|
|
scope_bodies.append((nid, node))
|
|
for child in node.children:
|
|
walk(child, nid)
|
|
return
|
|
|
|
if t == "module":
|
|
stmt = next((c for c in node.children if c.type == "module_statement"), None)
|
|
name = _fortran_name(stmt) if stmt else None
|
|
if name:
|
|
nid = _make_id(stem, name)
|
|
line = node.start_point[0] + 1
|
|
add_node(nid, name, line)
|
|
add_edge(file_nid, nid, "defines", line)
|
|
for child in node.children:
|
|
walk(child, nid)
|
|
return
|
|
|
|
# subroutines/functions inside a module live under internal_procedures
|
|
if t == "internal_procedures":
|
|
for child in node.children:
|
|
walk(child, scope_nid)
|
|
return
|
|
|
|
if t == "derived_type_definition":
|
|
stmt = next((c for c in node.children if c.type == "derived_type_statement"), None)
|
|
if stmt is not None:
|
|
name_node = next((c for c in stmt.children if c.type == "type_name"), None)
|
|
if name_node is not None:
|
|
type_name = _read_text(name_node, source).lower()
|
|
type_nid = _make_id(stem, type_name)
|
|
line = node.start_point[0] + 1
|
|
add_node(type_nid, type_name, line)
|
|
add_edge(scope_nid, type_nid, "defines", line)
|
|
return
|
|
|
|
if t == "subroutine":
|
|
stmt = next((c for c in node.children if c.type == "subroutine_statement"), None)
|
|
name = _fortran_name(stmt) if stmt else None
|
|
if name:
|
|
nid = _make_id(stem, name)
|
|
line = node.start_point[0] + 1
|
|
add_node(nid, f"{name}()", line)
|
|
add_edge(scope_nid, nid, "defines", line)
|
|
scope_bodies.append((nid, node))
|
|
emit_signature_refs(node, nid, is_function=False)
|
|
for child in node.children:
|
|
walk(child, nid)
|
|
return
|
|
|
|
if t == "function":
|
|
stmt = next((c for c in node.children if c.type == "function_statement"), None)
|
|
name = _fortran_name(stmt) if stmt else None
|
|
if name:
|
|
nid = _make_id(stem, name)
|
|
line = node.start_point[0] + 1
|
|
add_node(nid, f"{name}()", line)
|
|
add_edge(scope_nid, nid, "defines", line)
|
|
scope_bodies.append((nid, node))
|
|
emit_signature_refs(node, nid, is_function=True)
|
|
for child in node.children:
|
|
walk(child, nid)
|
|
return
|
|
|
|
if t == "use_statement":
|
|
line = node.start_point[0] + 1
|
|
# tree-sitter-fortran uses module_name node for the used module
|
|
name_node = next((c for c in node.children if c.type in ("module_name", "name", "identifier")), None)
|
|
if name_node:
|
|
mod_name = _read_text(name_node, source).lower()
|
|
imp_nid = _make_id(mod_name)
|
|
add_node(imp_nid, mod_name, line)
|
|
add_edge(scope_nid, imp_nid, "imports", line, context="use")
|
|
return
|
|
|
|
for child in node.children:
|
|
walk(child, scope_nid)
|
|
|
|
walk(root, file_nid)
|
|
|
|
_stmt_headers = {
|
|
"subroutine_statement", "function_statement",
|
|
"program_statement", "module_statement",
|
|
}
|
|
for scope_nid, body_node in scope_bodies:
|
|
for child in body_node.children:
|
|
if child.type not in _stmt_headers:
|
|
walk_calls(child, scope_nid)
|
|
|
|
return {"nodes": nodes, "edges": edges}
|