flux_middle/rty/
normalize.rs

1use std::ops::ControlFlow;
2
3use itertools::Itertools;
4use rustc_data_structures::{fx::FxIndexSet, unord::UnordMap};
5use rustc_hir::def_id::{CrateNum, DefIndex, LOCAL_CRATE};
6use rustc_macros::{TyDecodable, TyEncodable};
7use toposort_scc::IndexGraph;
8
9use super::{ESpan, fold::TypeSuperFoldable};
10use crate::{
11    def_id::{FluxDefId, FluxId, FluxLocalDefId},
12    global_env::GlobalEnv,
13    rty::{
14        Binder, Expr, ExprKind, SortArg,
15        expr::SpecFuncKind,
16        fold::{TypeFoldable, TypeFolder, TypeSuperVisitable, TypeVisitable, TypeVisitor},
17    },
18};
19
20#[derive(TyEncodable, TyDecodable)]
21pub struct NormalizedDefns {
22    krate: CrateNum,
23    inlined_bodies: UnordMap<FluxId<DefIndex>, Binder<Expr>>,
24    /// Information about all function definitions both with a body and UIF
25    info: UnordMap<FluxId<DefIndex>, FuncInfo>,
26}
27
28// This implementation is needed for `flux-metada::Tables`
29impl Default for NormalizedDefns {
30    fn default() -> Self {
31        Self { krate: LOCAL_CRATE, inlined_bodies: UnordMap::default(), info: UnordMap::default() }
32    }
33}
34
35// TODO(nilehmann) should we make an enum? most of the fields don't matter for UIFs
36/// This type represents what we know about a flux-def *after*
37/// normalization, i.e. after "inlining" all or some transitively
38/// called flux-defs.
39/// - When `FLUX_SMT_DEFINE_FUN=1` is set we inline
40///   all *polymorphic* flux-defs, since they cannot
41///   be represented  as `define-fun` in SMTLIB but leave
42///   all *monomorphic* flux-defs un-inlined.
43/// - When the above flag is not set, we replace *every* flux-def
44///   with its (transitively) inlined body
45#[derive(Clone, TyEncodable, TyDecodable)]
46pub struct FuncInfo {
47    /// Whether or not this function is inlined (i.e. NOT represented as `define-fun`).
48    /// This value is irrelevant of UIFs.
49    pub inline: bool,
50    /// Whether or not this function is uninterpreted by default
51    /// This value is irrelevant of UIFs.
52    pub hide: bool,
53    /// The rank of this function in the topological sort of all the flux-defs, needed so
54    /// we can specify the `define-fun` in the correct order, without any "forward"
55    /// dependencies which the SMT solver cannot handle.
56    pub rank: usize,
57    /// Whether the function is a UIF
58    pub uif: bool,
59}
60
61#[derive(Default)]
62pub(super) struct InliningCtxt {
63    inlined_bodies: UnordMap<FluxLocalDefId, Binder<Expr>>,
64    info: UnordMap<FluxLocalDefId, FuncInfo>,
65}
66
67pub(super) struct Normalizer<'a, 'genv, 'tcx> {
68    genv: GlobalEnv<'genv, 'tcx>,
69    inlining: Option<&'a InliningCtxt>,
70}
71
72impl NormalizedDefns {
73    pub fn new(
74        genv: GlobalEnv,
75        funs: &[(FluxLocalDefId, Option<Binder<Expr>>, bool)],
76    ) -> Result<Self, Vec<FluxLocalDefId>> {
77        // 1. Topologically sort the Defns
78        let ds = toposort(funs)?;
79
80        // 2. Expand each defn in the sorted order
81        let mut inlining = InliningCtxt::default();
82        for (rank, i) in ds.iter().enumerate() {
83            let (id, body, hide) = &funs[*i];
84
85            if let Some(body) = body {
86                let body = body.fold_with(&mut Normalizer::new(genv, Some(&inlining)));
87
88                inlining.inlined_bodies.insert(*id, body);
89                inlining.info.insert(
90                    *id,
91                    FuncInfo {
92                        rank,
93                        inline: genv.should_inline_fun(id.to_def_id()),
94                        hide: *hide,
95                        uif: false,
96                    },
97                );
98            } else {
99                inlining
100                    .info
101                    .insert(*id, FuncInfo { rank, inline: false, hide: *hide, uif: true });
102            }
103        }
104        Ok(Self {
105            krate: LOCAL_CRATE,
106            info: inlining
107                .info
108                .into_items()
109                .map(|(id, info)| (id.local_def_index(), info))
110                .collect(),
111            inlined_bodies: inlining
112                .inlined_bodies
113                .into_items()
114                .map(|(id, body)| (id.local_def_index(), body))
115                .collect(),
116        })
117    }
118
119    pub fn func_info(&self, did: FluxDefId) -> FuncInfo {
120        debug_assert_eq!(self.krate, did.krate());
121        self.info.get(&did.index()).unwrap().clone()
122    }
123
124    pub fn inlined_body(&self, did: FluxDefId) -> Binder<Expr> {
125        debug_assert_eq!(self.krate, did.krate());
126        self.inlined_bodies.get(&did.index()).unwrap().clone()
127    }
128}
129
130/// Returns
131/// * either Ok(d1...dn) which are topologically sorted such that
132///   forall i < j, di does not depend on i.e. "call" dj
133/// * or Err(d1...dn) where d1 'calls' d2 'calls' ... 'calls' dn 'calls' d1
134fn toposort<T>(
135    defns: &[(FluxLocalDefId, Option<Binder<Expr>>, T)],
136) -> Result<Vec<usize>, Vec<FluxLocalDefId>> {
137    // 1. Make a Symbol to Index map
138    let s2i: UnordMap<FluxLocalDefId, usize> = defns
139        .iter()
140        .enumerate()
141        .map(|(i, defn)| (defn.0, i))
142        .collect();
143
144    // 2. Make the dependency graph
145    let mut adj_list = Vec::with_capacity(defns.len());
146    for defn in defns {
147        if let Some(body) = &defn.1 {
148            let deps = local_deps(body)
149                .iter()
150                .filter_map(|s| s2i.get(s).copied())
151                .collect_vec();
152            adj_list.push(deps);
153        } else {
154            adj_list.push(vec![]);
155        }
156    }
157    let mut g = IndexGraph::from_adjacency_list(&adj_list);
158    g.transpose();
159
160    // 3. Topologically sort the graph
161    match g.toposort_or_scc() {
162        Ok(is) => Ok(is),
163        Err(mut scc) => {
164            let cycle = scc.pop().unwrap();
165            Err(cycle.iter().map(|i| defns[*i].0).collect())
166        }
167    }
168}
169
170pub fn local_deps(body: &Binder<Expr>) -> FxIndexSet<FluxLocalDefId> {
171    struct DepsVisitor(FxIndexSet<FluxLocalDefId>);
172    impl TypeVisitor for DepsVisitor {
173        #[allow(clippy::disallowed_methods, reason = "refinement functions cannot be extern specs")]
174        fn visit_expr(&mut self, expr: &Expr) -> ControlFlow<!> {
175            if let ExprKind::App(func, ..) = expr.kind()
176                && let ExprKind::GlobalFunc(SpecFuncKind::Def(did)) = func.kind()
177                && let Some(did) = did.as_local()
178            {
179                self.0.insert(did);
180            }
181            expr.super_visit_with(self)
182        }
183    }
184    let mut visitor = DepsVisitor(Default::default());
185    let _ = body.visit_with(&mut visitor);
186    visitor.0
187}
188
189impl<'a, 'genv, 'tcx> Normalizer<'a, 'genv, 'tcx> {
190    pub(super) fn new(genv: GlobalEnv<'genv, 'tcx>, inlining: Option<&'a InliningCtxt>) -> Self {
191        Self { genv, inlining }
192    }
193
194    #[allow(clippy::disallowed_methods, reason = "refinement functions cannot be extern specs")]
195    fn func_defn(&self, did: FluxDefId) -> Binder<Expr> {
196        if let Some(inlining) = self.inlining
197            && let Some(local_id) = did.as_local()
198        {
199            inlining.inlined_bodies[&local_id].clone()
200        } else {
201            self.genv.inlined_body(did)
202        }
203    }
204
205    #[allow(clippy::disallowed_methods, reason = "refinement functions cannot be extern specs")]
206    fn should_inline(&self, did: FluxDefId) -> bool {
207        let info = if let Some(inlining) = self.inlining
208            && let Some(local_id) = did.as_local()
209        {
210            &inlining.info[&local_id]
211        } else {
212            &self.genv.normalized_info(did)
213        };
214        info.inline && !info.hide
215    }
216
217    fn at_base(expr: Expr, espan: Option<ESpan>) -> Expr {
218        match espan {
219            Some(espan) => BaseSpanner::new(espan).fold_expr(&expr),
220            None => expr,
221        }
222    }
223
224    fn app(
225        &mut self,
226        func: &Expr,
227        sort_args: &[SortArg],
228        args: &[Expr],
229        espan: Option<ESpan>,
230    ) -> Expr {
231        match func.kind() {
232            ExprKind::GlobalFunc(SpecFuncKind::Def(did)) if self.should_inline(*did) => {
233                let res = self.func_defn(*did).replace_bound_refts(args);
234                Self::at_base(res, espan)
235            }
236            ExprKind::Abs(lam) => {
237                let res = lam.apply(args);
238                Self::at_base(res, espan)
239            }
240            _ => Expr::app(func.clone(), sort_args.into(), args.into()).at_opt(espan),
241        }
242    }
243}
244
245impl TypeFolder for Normalizer<'_, '_, '_> {
246    fn fold_expr(&mut self, expr: &Expr) -> Expr {
247        let expr = expr.super_fold_with(self);
248        let span = expr.span();
249        match expr.kind() {
250            ExprKind::App(func, sorts, args) => self.app(func, sorts, args, span),
251            ExprKind::FieldProj(e, proj) => e.proj_and_reduce(*proj),
252            _ => expr,
253        }
254    }
255}
256
257struct BaseSpanner {
258    espan: ESpan,
259}
260
261impl BaseSpanner {
262    fn new(espan: ESpan) -> Self {
263        Self { espan }
264    }
265}
266
267impl TypeFolder for BaseSpanner {
268    fn fold_expr(&mut self, expr: &Expr) -> Expr {
269        expr.super_fold_with(self).at_base(self.espan)
270    }
271}
flux_middle/rty/normalize.rs

flux_middle/rty/
normalize.rs
