use core::fmt;
use crate::frontend::{Lexeme, Token, TokenKind};
use crate::ir::{ByteRange, IrNode, SUMMARY_HEAD_TOKENS, Shape, SyntaxIrFile};
pub const FEATURE_SCHEMA_VERSION: &str = "ir-features-v1";
pub const WINDOW_LENGTHS: &[usize] = &[4, 8, 16];
pub const MIN_SUBTREE_NODES: usize = 5;
pub const SHAPE_TAG_SLOTS: usize = 23;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum FeatureKind {
StatementWindow,
Subtree,
Cfg,
ApiCallSequence,
ApiCallMultiset,
}
impl FeatureKind {
pub const ALL: [Self; 5] = [
Self::StatementWindow,
Self::Subtree,
Self::Cfg,
Self::ApiCallSequence,
Self::ApiCallMultiset,
];
#[must_use]
pub const fn name(self) -> &'static str {
match self {
Self::StatementWindow => "statement_window",
Self::Subtree => "subtree",
Self::Cfg => "cfg",
Self::ApiCallSequence => "api_call_sequence",
Self::ApiCallMultiset => "api_call_multiset",
}
}
#[must_use]
pub fn from_name(name: &str) -> Option<Self> {
Self::ALL.into_iter().find(|kind| kind.name() == name)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct FeatureHash([u8; 16]);
impl FeatureHash {
#[must_use]
pub const fn from_bytes(bytes: [u8; 16]) -> Self {
Self(bytes)
}
#[must_use]
pub const fn as_bytes(&self) -> &[u8; 16] {
&self.0
}
#[must_use]
pub fn to_hex(&self) -> String {
self.to_string()
}
}
impl fmt::Display for FeatureHash {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
for byte in self.0 {
write!(f, "{byte:02x}")?;
}
Ok(())
}
}
struct FeatureHasher {
hasher: blake3::Hasher,
}
impl FeatureHasher {
fn new(domain: &str) -> Self {
let mut this = Self {
hasher: blake3::Hasher::new(),
};
this.write_bytes(domain.as_bytes());
this.write_bytes(FEATURE_SCHEMA_VERSION.as_bytes());
this
}
fn write_bytes(&mut self, bytes: &[u8]) {
let len = u32::try_from(bytes.len()).unwrap_or(u32::MAX);
self.hasher.update(&len.to_le_bytes());
self.hasher.update(bytes);
}
fn write_str(&mut self, text: &str) {
self.write_bytes(text.as_bytes());
}
fn write_u8(&mut self, value: u8) {
self.hasher.update(&[value]);
}
fn write_u32(&mut self, value: u32) {
self.hasher.update(&value.to_le_bytes());
}
fn finish(self) -> FeatureHash {
let digest = self.hasher.finalize();
let mut out = [0u8; 16];
out.copy_from_slice(&digest.as_bytes()[..16]);
FeatureHash(out)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FileFeatures {
pub units: Vec<UnitFeatures>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UnitFeatures {
pub name: Option<Lexeme>,
pub shape_tag: u8,
pub range: ByteRange,
pub windows: Vec<WindowFeature>,
pub subtrees: Vec<SubtreeFeature>,
pub vector: CharacteristicVector,
pub cfg: CfgFeature,
pub api: ApiCallFeature,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct UnitRef {
pub file: usize,
pub unit: usize,
pub node_count: u32,
}
impl UnitRef {
#[must_use]
pub fn within_length_ratio(self, other: Self, max_ratio: f64) -> bool {
let (small, large) = if self.node_count <= other.node_count {
(self.node_count, other.node_count)
} else {
(other.node_count, self.node_count)
};
if small == 0 {
return large == 0;
}
f64::from(large) / f64::from(small) <= max_ratio
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WindowFeature {
pub hash: FeatureHash,
pub length: usize,
pub range: ByteRange,
pub block: u32,
pub offset: u32,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SubtreeFeature {
pub hash: FeatureHash,
pub node_count: usize,
pub range: ByteRange,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CharacteristicVector {
pub counts: [u32; SHAPE_TAG_SLOTS],
pub max_depth: u32,
pub node_count: u32,
}
impl CharacteristicVector {
#[must_use]
pub fn l1_distance(&self, other: &Self) -> u64 {
self.counts
.iter()
.zip(other.counts.iter())
.map(|(&a, &b)| u64::from(a.abs_diff(b)))
.sum()
}
#[must_use]
pub fn shape_divergence(&self, other: &Self) -> f64 {
let span = u64::from(self.node_count) + u64::from(other.node_count);
if span == 0 {
return 0.0;
}
#[expect(
clippy::cast_precision_loss,
reason = "node counts of this size lose nothing a threshold comparison would notice"
)]
{
self.l1_distance(other) as f64 / span as f64
}
}
#[must_use]
pub fn cosine_similarity(&self, other: &Self) -> f64 {
if self.counts.iter().all(|&c| c == 0) || other.counts.iter().all(|&c| c == 0) {
return 0.0;
}
let mut dot = 0.0f64;
let mut norm_self = 0.0f64;
let mut norm_other = 0.0f64;
for (&a, &b) in self.counts.iter().zip(other.counts.iter()) {
let (a, b) = (f64::from(a), f64::from(b));
dot = a.mul_add(b, dot);
norm_self = a.mul_add(a, norm_self);
norm_other = b.mul_add(b, norm_other);
}
dot / (norm_self * norm_other).sqrt()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CfgFeature {
pub hash: FeatureHash,
pub skeleton_hash: FeatureHash,
pub op_count: u32,
pub skeleton_ops: u32,
pub max_loop_depth: u32,
pub branch_count: u32,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ApiCallFeature {
pub names: Vec<Lexeme>,
pub sequence_hash: FeatureHash,
pub multiset_hash: FeatureHash,
}
const OP_LOOP_ENTER: u8 = 1;
const OP_LOOP_EXIT: u8 = 2;
const OP_BRANCH_ENTER: u8 = 3;
const OP_BRANCH_EXIT: u8 = 4;
const OP_MATCH_ENTER: u8 = 5;
const OP_MATCH_EXIT: u8 = 6;
const OP_ARM_ENTER: u8 = 7;
const OP_ARM_EXIT: u8 = 8;
const OP_TRY: u8 = 9;
const OP_RETURN: u8 = 10;
const OP_BREAK: u8 = 11;
const OP_CONTINUE: u8 = 12;
const OP_CALL: u8 = 13;
#[must_use]
pub fn extract(file: &SyntaxIrFile) -> FileFeatures {
let mut units = Vec::new();
file.walk(&mut |node| {
if matches!(node.shape, Shape::Function | Shape::Method | Shape::Closure) {
units.push(unit_features(node, &file.tokens));
}
});
FileFeatures { units }
}
fn unit_features(unit: &IrNode, tokens: &[Token]) -> UnitFeatures {
let mut windows = Vec::new();
let mut block = 0u32;
unit.walk(&mut |node| {
if matches!(node.shape, Shape::Block) {
block_windows(node, block, tokens, &mut windows);
block = block.saturating_add(1);
}
});
let mut subtrees = Vec::new();
let _ = subtree_features(unit, &mut subtrees);
let mut vector = CharacteristicVector::default();
accumulate_vector(unit, 1, &mut vector);
UnitFeatures {
name: unit.name.clone(),
shape_tag: unit.shape.tag(),
range: unit.range,
windows,
subtrees,
vector,
cfg: cfg_feature(unit),
api: api_feature(unit, tokens),
}
}
fn native_kind(shape: &Shape) -> &str {
match shape {
Shape::Native(kind) => kind.as_str(),
_ => "",
}
}
fn block_windows(block: &IrNode, ordinal: u32, tokens: &[Token], out: &mut Vec<WindowFeature>) {
let statements: Vec<&IrNode> = block
.children
.iter()
.filter(|child| child.shape.is_statement() || matches!(child.shape, Shape::Native(_)))
.collect();
for &length in WINDOW_LENGTHS {
for (offset, window) in statements.windows(length).enumerate() {
let mut hasher = FeatureHasher::new("stmt-window");
hasher.write_u32(u32::try_from(length).unwrap_or(u32::MAX));
for statement in window {
write_statement(&mut hasher, statement, tokens);
}
out.push(WindowFeature {
hash: hasher.finish(),
length,
range: ByteRange {
start: window[0].range.start,
end: window[length - 1].range.end,
},
block: ordinal,
offset: u32::try_from(offset).unwrap_or(u32::MAX),
});
}
}
}
fn write_statement(hasher: &mut FeatureHasher, statement: &IrNode, tokens: &[Token]) {
hasher.write_u8(statement.shape.tag());
hasher.write_str(native_kind(&statement.shape));
let end = statement.token_end.min(tokens.len());
let start = statement.token_start.min(end);
let head_tags: Vec<u8> = tokens[start..end]
.iter()
.take(SUMMARY_HEAD_TOKENS)
.map(|token| token.kind.tag())
.collect();
hasher.write_bytes(&head_tags);
}
fn subtree_features(node: &IrNode, out: &mut Vec<SubtreeFeature>) -> (FeatureHash, usize) {
struct Frame<'a> {
node: &'a IrNode,
next_child: usize,
child_hashes: Vec<FeatureHash>,
node_count: usize,
}
let mut pending = vec![Frame {
node,
next_child: 0,
child_hashes: Vec::with_capacity(node.children.len()),
node_count: 1,
}];
loop {
let Some(frame) = pending.last_mut() else {
unreachable!("the root frame is retained until its result is returned");
};
if let Some(child) = frame.node.children.get(frame.next_child) {
frame.next_child += 1;
pending.push(Frame {
node: child,
next_child: 0,
child_hashes: Vec::with_capacity(child.children.len()),
node_count: 1,
});
continue;
}
let mut hasher = FeatureHasher::new("subtree");
hasher.write_u8(frame.node.shape.tag());
hasher.write_str(native_kind(&frame.node.shape));
hasher.write_u32(u32::try_from(frame.child_hashes.len()).unwrap_or(u32::MAX));
for hash in &frame.child_hashes {
hasher.write_bytes(hash.as_bytes());
}
let hash = hasher.finish();
if frame.node_count >= MIN_SUBTREE_NODES {
out.push(SubtreeFeature {
hash,
node_count: frame.node_count,
range: frame.node.range,
});
}
let result = (hash, frame.node_count);
pending.pop();
let Some(parent) = pending.last_mut() else {
return result;
};
parent.child_hashes.push(result.0);
parent.node_count += result.1;
}
}
fn accumulate_vector(node: &IrNode, depth: u32, vector: &mut CharacteristicVector) {
let mut pending = vec![(node, depth)];
while let Some((node, depth)) = pending.pop() {
vector.node_count += 1;
vector.max_depth = vector.max_depth.max(depth);
vector.counts[usize::from(node.shape.tag())] += 1;
pending.extend(
node.children
.iter()
.rev()
.map(|child| (child, depth.saturating_add(1))),
);
}
}
#[derive(Default)]
struct CfgWalk {
ops: Vec<u8>,
skeleton: Vec<u8>,
op_count: u32,
skeleton_ops: u32,
branch_count: u32,
loop_depth: u32,
max_loop_depth: u32,
}
impl CfgWalk {
fn push_op(&mut self, op: u8) {
self.ops.push(op);
self.op_count += 1;
if op != OP_CALL {
self.skeleton.push(op);
self.skeleton_ops += 1;
}
}
fn push_operand(&mut self, bytes: &[u8]) {
self.ops.extend_from_slice(bytes);
self.skeleton.extend_from_slice(bytes);
}
fn visit(&mut self, node: &IrNode) {
enum Visit<'a> {
Enter(&'a IrNode),
Exit { op: u8, leaves_loop: bool },
}
let mut pending = vec![Visit::Enter(node)];
while let Some(visit) = pending.pop() {
match visit {
Visit::Exit { op, leaves_loop } => {
if leaves_loop {
self.loop_depth -= 1;
}
self.push_op(op);
}
Visit::Enter(node) => {
let exit = match &node.shape {
Shape::Loop => {
self.push_op(OP_LOOP_ENTER);
self.loop_depth += 1;
self.max_loop_depth = self.max_loop_depth.max(self.loop_depth);
Some(Visit::Exit {
op: OP_LOOP_EXIT,
leaves_loop: true,
})
}
Shape::Branch => {
self.push_op(OP_BRANCH_ENTER);
self.branch_count += 1;
Some(Visit::Exit {
op: OP_BRANCH_EXIT,
leaves_loop: false,
})
}
Shape::Match => {
self.push_op(OP_MATCH_ENTER);
let arms = node
.children
.iter()
.filter(|child| matches!(child.shape, Shape::MatchArm))
.count();
let arms = u32::try_from(arms).unwrap_or(u32::MAX);
self.push_operand(&arms.to_le_bytes());
Some(Visit::Exit {
op: OP_MATCH_EXIT,
leaves_loop: false,
})
}
Shape::MatchArm => {
self.push_op(OP_ARM_ENTER);
Some(Visit::Exit {
op: OP_ARM_EXIT,
leaves_loop: false,
})
}
Shape::Try => {
self.push_op(OP_TRY);
None
}
Shape::Return => {
self.push_op(OP_RETURN);
None
}
Shape::Break => {
self.push_op(OP_BREAK);
None
}
Shape::Continue => {
self.push_op(OP_CONTINUE);
None
}
Shape::Call => {
self.push_op(OP_CALL);
None
}
_ => None,
};
if let Some(exit) = exit {
pending.push(exit);
}
pending.extend(node.children.iter().rev().map(Visit::Enter));
}
}
}
}
}
fn cfg_feature(unit: &IrNode) -> CfgFeature {
let mut walk = CfgWalk::default();
walk.visit(unit);
let mut hasher = FeatureHasher::new("cfg");
hasher.write_bytes(&walk.ops);
let mut skeleton = FeatureHasher::new("cfg-skeleton");
skeleton.write_bytes(&walk.skeleton);
CfgFeature {
hash: hasher.finish(),
skeleton_hash: skeleton.finish(),
op_count: walk.op_count,
skeleton_ops: walk.skeleton_ops,
max_loop_depth: walk.max_loop_depth,
branch_count: walk.branch_count,
}
}
fn callee_name(call: &IrNode, tokens: &[Token]) -> Option<Lexeme> {
let end = call.token_end.min(tokens.len());
let start = call.token_start.min(end);
let slice = &tokens[start..end];
let open = slice
.iter()
.position(|token| matches!(token.kind, TokenKind::Punctuation) && token.text == "(")?;
slice[..open]
.iter()
.rev()
.find(|token| matches!(token.kind, TokenKind::Identifier))
.map(|token| token.text.clone())
}
fn api_feature(unit: &IrNode, tokens: &[Token]) -> ApiCallFeature {
let mut names: Vec<Lexeme> = Vec::new();
unit.walk(&mut |node| {
if matches!(node.shape, Shape::Call) {
if let Some(name) = callee_name(node, tokens) {
names.push(name);
}
}
});
let mut sequence = FeatureHasher::new("api-call");
for name in &names {
sequence.write_str(name);
}
let mut sorted: Vec<&Lexeme> = names.iter().collect();
sorted.sort_unstable_by(|a, b| a.as_str().cmp(b.as_str()));
let mut multiset = FeatureHasher::new("api-call-set");
for name in sorted {
multiset.write_str(name);
}
ApiCallFeature {
names,
sequence_hash: sequence.finish(),
multiset_hash: multiset.finish(),
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests;