flux_middle/rty/
spec_func.rs

1use std::ops::ControlFlow;
2
3use itertools::Itertools;
4use rustc_data_structures::{
5    fx::{FxIndexMap, FxIndexSet},
6    unord::UnordMap,
7};
8use rustc_hir::def_id::DefIndex;
9use rustc_macros::{TyDecodable, TyEncodable};
10use toposort_scc::IndexGraph;
11
12use crate::{
13    def_id::{FluxDefId, FluxId, FluxLocalDefId},
14    rty::{
15        Binder, Expr, ExprKind,
16        expr::SpecFuncKind,
17        fold::{TypeSuperVisitable, TypeVisitable, TypeVisitor},
18    },
19};
20
21/// A spec function (aka flux-def) as written by the user, i.e., *before* inlining any of the
22/// functions it calls.
23#[derive(Clone, TyEncodable, TyDecodable)]
24pub enum SpecFunc {
25    /// An uninterpreted function
26    Uif,
27    /// A function with a body
28    Defined {
29        body: Binder<Expr>,
30        /// Whether the function is uninterpreted by default
31        hide: bool,
32    },
33}
34
35impl SpecFunc {
36    pub fn is_uif(&self) -> bool {
37        matches!(self, SpecFunc::Uif)
38    }
39}
40
41/// All the spec functions of a crate sorted topologically, i.e., every function comes after all
42/// the functions (in the same crate) it calls.
43///
44/// Functions are identified by their index within the crate. Use [`GlobalEnv::spec_funcs`] to
45/// get the spec functions of the crate a [`FluxDefId`] belongs to.
46///
47/// [`GlobalEnv::spec_funcs`]: crate::global_env::GlobalEnv::spec_funcs
48#[derive(Default, TyEncodable, TyDecodable)]
49pub struct SpecFuncs {
50    funcs: FxIndexMap<FluxId<DefIndex>, SpecFunc>,
51}
52
53impl SpecFuncs {
54    /// Sorts the spec functions of the local crate topologically. Returns `Err(cycle)` if there's
55    /// a cycle in the functions.
56    pub fn new(funcs: Vec<(FluxLocalDefId, SpecFunc)>) -> Result<Self, Vec<FluxLocalDefId>> {
57        let order = toposort(&funcs)?;
58        let mut funcs = funcs.into_iter().map(Some).collect_vec();
59        let funcs = order
60            .into_iter()
61            .map(|i| {
62                let (id, func) = funcs[i].take().unwrap();
63                (id.local_def_index(), func)
64            })
65            .collect();
66        Ok(Self { funcs })
67    }
68
69    pub fn get(&self, id: FluxId<DefIndex>) -> &SpecFunc {
70        &self.funcs[&id]
71    }
72
73    /// The position of the function in the topological order. A function can only call functions
74    /// in the same crate with a lower rank.
75    pub fn rank(&self, id: FluxId<DefIndex>) -> usize {
76        self.funcs.get_index_of(&id).unwrap()
77    }
78}
79
80/// Returns
81/// * either Ok(d1...dn) which are topologically sorted such that
82///   forall i < j, di does not depend on i.e. "call" dj
83/// * or Err(d1...dn) where d1 'calls' d2 'calls' ... 'calls' dn 'calls' d1
84fn toposort(funcs: &[(FluxLocalDefId, SpecFunc)]) -> Result<Vec<usize>, Vec<FluxLocalDefId>> {
85    // 1. Make a Symbol to Index map
86    let s2i: UnordMap<FluxDefId, usize> = funcs
87        .iter()
88        .enumerate()
89        .map(|(i, (id, _))| (id.to_def_id(), i))
90        .collect();
91
92    // 2. Make the dependency graph. Calls to functions in other crates are not in `s2i`.
93    let mut adj_list = Vec::with_capacity(funcs.len());
94    for (_, func) in funcs {
95        if let SpecFunc::Defined { body, .. } = func {
96            let deps = deps(body)
97                .iter()
98                .filter_map(|s| s2i.get(s).copied())
99                .collect_vec();
100            adj_list.push(deps);
101        } else {
102            adj_list.push(vec![]);
103        }
104    }
105    let mut g = IndexGraph::from_adjacency_list(&adj_list);
106    g.transpose();
107
108    // 3. Topologically sort the graph
109    match g.toposort_or_scc() {
110        Ok(is) => Ok(is),
111        Err(mut scc) => {
112            let cycle = scc.pop().unwrap();
113            Err(cycle.iter().map(|i| funcs[*i].0).collect())
114        }
115    }
116}
117
118/// Returns all the spec functions (of any crate) called in `body`
119pub fn deps(body: &Binder<Expr>) -> FxIndexSet<FluxDefId> {
120    struct DepsVisitor(FxIndexSet<FluxDefId>);
121    impl TypeVisitor for DepsVisitor {
122        fn visit_expr(&mut self, expr: &Expr) -> ControlFlow<!> {
123            if let ExprKind::App(func, ..) = expr.kind()
124                && let ExprKind::GlobalFunc(SpecFuncKind::Def(did)) = func.kind()
125            {
126                self.0.insert(*did);
127            }
128            expr.super_visit_with(self)
129        }
130    }
131    let mut visitor = DepsVisitor(Default::default());
132    let _ = body.visit_with(&mut visitor);
133    visitor.0
134}
flux_middle/rty/spec_func.rs

flux_middle/rty/
spec_func.rs
