use petgraph::{
algo::tarjan_scc,
graph::{EdgeIndex, NodeIndex},
visit::EdgeRef,
Graph,
};
use solana_libra_types::vm_error::{StatusCode, VMStatus};
use solana_libra_vm::{
access::ModuleAccess,
file_format::{
Bytecode, CompiledModule, FunctionDefinition, FunctionDefinitionIndex, FunctionHandleIndex,
LocalsSignatureIndex, SignatureToken, TypeParameterIndex,
},
};
use std::collections::{hash_map, HashMap, HashSet};
#[derive(Eq, PartialEq, Hash, Copy, Clone)]
struct Node(FunctionDefinitionIndex, TypeParameterIndex);
enum Edge<'a> {
Identity,
TyConApp(&'a SignatureToken),
}
pub struct InstantiationLoopChecker<'a> {
module: &'a CompiledModule,
graph: Graph<Node, Edge<'a>>,
node_map: HashMap<Node, NodeIndex>,
func_handle_def_map: HashMap<FunctionHandleIndex, FunctionDefinitionIndex>,
}
impl<'a> InstantiationLoopChecker<'a> {
pub fn new(module: &'a CompiledModule) -> Self {
Self {
module,
graph: Graph::new(),
node_map: HashMap::new(),
func_handle_def_map: module
.function_defs()
.iter()
.enumerate()
.map(|(def_idx, def)| (def.function, FunctionDefinitionIndex::new(def_idx as u16)))
.collect(),
}
}
fn get_or_add_node(&mut self, node: Node) -> NodeIndex {
match self.node_map.entry(node) {
hash_map::Entry::Occupied(entry) => *entry.get(),
hash_map::Entry::Vacant(entry) => {
let idx = self.graph.add_node(node);
entry.insert(idx);
idx
}
}
}
fn extract_type_parameters(ty: &SignatureToken) -> HashSet<TypeParameterIndex> {
use SignatureToken::*;
let mut type_params = HashSet::new();
fn rec(type_params: &mut HashSet<TypeParameterIndex>, ty: &SignatureToken) {
match ty {
Bool | Address | U64 | String | ByteArray => (),
TypeParameter(idx) => {
type_params.insert(*idx);
}
Reference(ty) | MutableReference(ty) => rec(type_params, ty),
Struct(_, tys) => {
for ty in tys {
rec(type_params, ty);
}
}
}
}
rec(&mut type_params, ty);
type_params
}
fn add_edge(&mut self, node_from: Node, node_to: Node, edge: Edge<'a>) {
let node_from_idx = self.get_or_add_node(node_from);
let node_to_idx = self.get_or_add_node(node_to);
self.graph.add_edge(node_from_idx, node_to_idx, edge);
}
fn build_graph_call(
&mut self,
caller_idx: FunctionDefinitionIndex,
callee_idx: FunctionDefinitionIndex,
type_actuals_idx: LocalsSignatureIndex,
) {
let type_actuals = &self.module.locals_signature_at(type_actuals_idx).0;
for (formal_idx, ty) in type_actuals.iter().enumerate() {
let formal_idx = formal_idx as TypeParameterIndex;
match ty {
SignatureToken::TypeParameter(actual_idx) => self.add_edge(
Node(caller_idx, *actual_idx),
Node(callee_idx, formal_idx),
Edge::Identity,
),
_ => {
for type_param in Self::extract_type_parameters(ty) {
self.add_edge(
Node(caller_idx, type_param),
Node(callee_idx, formal_idx),
Edge::TyConApp(&ty),
);
}
}
}
}
}
fn build_graph_function_def(
&mut self,
caller_idx: FunctionDefinitionIndex,
caller_def: &FunctionDefinition,
) {
for instr in &caller_def.code.code {
if let Bytecode::Call(callee_handle_idx, type_actuals_idx) = instr {
if let Some(callee_idx) = self.func_handle_def_map.get(&callee_handle_idx) {
let callee_idx = *callee_idx;
self.build_graph_call(caller_idx, callee_idx, *type_actuals_idx)
}
}
}
}
fn build_graph(&mut self) {
for (def_idx, func_def) in self
.module
.function_defs()
.iter()
.filter(|def| !def.is_native())
.enumerate()
{
self.build_graph_function_def(FunctionDefinitionIndex::new(def_idx as u16), func_def)
}
}
fn find_non_trivial_components(&self) -> Vec<(Vec<NodeIndex>, Vec<EdgeIndex>)> {
tarjan_scc(&self.graph)
.into_iter()
.filter_map(move |nodes| {
let node_set: HashSet<_> = nodes.iter().cloned().collect();
let edges: Vec<_> = nodes
.iter()
.flat_map(|node_idx| {
self.graph.edges(*node_idx).filter_map(|edge| {
if node_set.contains(&edge.target()) {
Some(edge.id())
} else {
None
}
})
})
.collect();
if edges.iter().any(
|edge_idx| match self.graph.edge_weight(*edge_idx).unwrap() {
Edge::Identity => false,
Edge::TyConApp(_) => true,
},
) {
Some((nodes, edges))
} else {
None
}
})
.collect()
}
fn format_node(&self, node_idx: NodeIndex) -> String {
let Node(def_idx, param_idx) = self.graph.node_weight(node_idx).unwrap();
format!("f{}#{}", def_idx, param_idx)
}
fn format_edge(&self, edge_idx: EdgeIndex) -> String {
let (node_idx_1, node_idx_2) = self.graph.edge_endpoints(edge_idx).unwrap();
let node_1 = self.format_node(node_idx_1);
let node_2 = self.format_node(node_idx_2);
match self.graph.edge_weight(edge_idx).unwrap() {
Edge::TyConApp(ty) => format!("{} --{:?}--> {}", node_1, ty, node_2,),
Edge::Identity => format!("{} ----> {}", node_1, node_2),
}
}
pub fn verify(mut self) -> Vec<VMStatus> {
self.build_graph();
let components = self.find_non_trivial_components();
components
.into_iter()
.map(|(nodes, edges)| {
let msg_edges = edges
.into_iter()
.filter_map(|edge_idx| match self.graph.edge_weight(edge_idx).unwrap() {
Edge::TyConApp(_) => Some(self.format_edge(edge_idx)),
_ => None,
})
.collect::<Vec<_>>()
.join(", ");
let msg_nodes = nodes
.into_iter()
.map(|node_idx| self.format_node(node_idx))
.collect::<Vec<_>>()
.join(", ");
let msg = format!(
"edges with constructors: [{}], nodes: [{}]",
msg_edges, msg_nodes
);
VMStatus::new(StatusCode::LOOP_IN_INSTANTIATION_GRAPH).with_message(msg)
})
.collect()
}
}