aprender-core 0.31.2

Next-generation machine learning library in pure Rust
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102


impl Default for GNNErrorEncoder {
    fn default() -> Self {
        Self::default_config()
    }
}

/// Program-feedback graph structure.
///
/// Per Yasunaga & Liang (2020), this graph connects symbols
/// in source code with diagnostic feedback for GNN reasoning.
#[derive(Debug, Clone)]
pub struct ProgramFeedbackGraph {
    /// Node features
    pub node_features: Vec<Vec<f32>>,
    /// Node types
    pub node_types: Vec<NodeType>,
    /// Edges (source, target)
    pub edges: Vec<(usize, usize)>,
    /// Edge types
    pub edge_types: Vec<EdgeType>,
}

impl ProgramFeedbackGraph {
    /// Create a new empty graph.
    #[must_use]
    pub fn new() -> Self {
        Self {
            node_features: Vec::new(),
            node_types: Vec::new(),
            edges: Vec::new(),
            edge_types: Vec::new(),
        }
    }

    /// Add a node to the graph.
    pub fn add_node(&mut self, node_type: NodeType, features: Vec<f32>) -> usize {
        let idx = self.node_features.len();
        self.node_features.push(features);
        self.node_types.push(node_type);
        idx
    }

    /// Add an edge to the graph.
    pub fn add_edge(&mut self, source: usize, target: usize, edge_type: EdgeType) {
        self.edges.push((source, target));
        self.edge_types.push(edge_type);
    }

    /// Get the number of nodes.
    #[must_use]
    pub fn num_nodes(&self) -> usize {
        self.node_features.len()
    }

    /// Get the number of edges.
    #[must_use]
    pub fn num_edges(&self) -> usize {
        self.edges.len()
    }
}

impl Default for ProgramFeedbackGraph {
    fn default() -> Self {
        Self::new()
    }
}

/// Node type in program-feedback graph.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NodeType {
    /// AST node (variable, type, expression, etc.)
    Ast,
    /// Compiler diagnostic
    Diagnostic,
    /// Expected type in type error
    ExpectedType,
    /// Found type in type error
    FoundType,
    /// Compiler suggestion
    Suggestion,
}

/// Edge type in program-feedback graph.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EdgeType {
    /// AST parent-child relationship
    AstChild,
    /// Data flow edge
    DataFlow,
    /// Control flow edge
    ControlFlow,
    /// Diagnostic refers to code location
    DiagnosticRefers,
    /// Type expectation
    Expects,
    /// Type found
    Found,
}

#[cfg(test)]
mod tests;