Skip to content

Commit 809279c

Browse files
coord-ecoord-e
committed
fixup! Translate and register formula_fn
1 parent 2d78295 commit 809279c

1 file changed

Lines changed: 296 additions & 0 deletions

File tree

‎src/analyze/annot_fn.rs‎

Lines changed: 296 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,296 @@
1+
use std::collections::HashMap;
2+
3+
use pretty::{termcolor, Pretty};
4+
use rustc_hir::{def_id::LocalDefId, HirId};
5+
use rustc_index::IndexVec;
6+
use rustc_middle::ty::{self as mir_ty, TyCtxt};
7+
8+
use crate::annot::AnnotFormula;
9+
use crate::chc;
10+
use crate::rty;
11+
12+
#[derive(Debug, Clone)]
13+
pub struct FormulaFn<'tcx> {
14+
params: IndexVec<rty::FunctionParamIdx, mir_ty::Ty<'tcx>>,
15+
formula: chc::Formula<rty::FunctionParamIdx>,
16+
}
17+
18+
impl<'a, 'b, D> Pretty<'a, D, termcolor::ColorSpec> for &'b FormulaFn<'_>
19+
where
20+
D: pretty::DocAllocator<'a, termcolor::ColorSpec>,
21+
D::Doc: Clone,
22+
{
23+
fn pretty(self, allocator: &'a D) -> pretty::DocBuilder<'a, D, termcolor::ColorSpec> {
24+
allocator
25+
.intersperse(
26+
self.params.iter_enumerated().map(|(idx, ty)| {
27+
idx.pretty(allocator)
28+
.append(": ")
29+
.append(allocator.as_string(ty))
30+
}),
31+
", ",
32+
)
33+
.enclose("|", "|")
34+
.group()
35+
.append(self.formula.pretty(allocator))
36+
.group()
37+
}
38+
}
39+
40+
impl<'tcx> FormulaFn<'tcx> {
41+
pub fn to_require_annot(
42+
&self,
43+
sig: mir_ty::FnSig<'tcx>,
44+
) -> Option<AnnotFormula<rty::FunctionParamIdx>> {
45+
if &self.params.raw != sig.inputs() {
46+
return None;
47+
}
48+
Some(AnnotFormula::Formula(self.formula.clone()))
49+
}
50+
51+
pub fn to_ensure_annot(
52+
&self,
53+
sig: mir_ty::FnSig<'tcx>,
54+
) -> Option<AnnotFormula<rty::RefinedTypeVar<rty::FunctionParamIdx>>> {
55+
if &self.params.raw[1..] != sig.inputs() {
56+
return None;
57+
}
58+
if self.params.raw[0] != sig.output() {
59+
return None;
60+
}
61+
Some(AnnotFormula::Formula(self.formula.clone().map_var(|v| {
62+
if v.as_usize() == 0 {
63+
rty::RefinedTypeVar::Value
64+
} else {
65+
rty::RefinedTypeVar::Free(rty::FunctionParamIdx::from(v.as_usize() - 1))
66+
}
67+
})))
68+
}
69+
}
70+
71+
#[derive(Debug, Clone, Copy)]
72+
enum AmbiguousBinOp {
73+
Eq,
74+
Ne,
75+
Ge,
76+
Le,
77+
Gt,
78+
Lt,
79+
}
80+
81+
#[derive(Debug, Clone)]
82+
enum FormulaOrTerm<T> {
83+
Formula(chc::Formula<T>),
84+
Term(chc::Term<T>),
85+
BinOp(chc::Term<T>, AmbiguousBinOp, chc::Term<T>),
86+
Not(Box<FormulaOrTerm<T>>),
87+
Literal(bool),
88+
}
89+
90+
impl<T> FormulaOrTerm<T> {
91+
fn into_formula(self) -> Option<chc::Formula<T>> {
92+
let fo = match self {
93+
FormulaOrTerm::Formula(fo) => fo,
94+
FormulaOrTerm::Term { .. } => return None,
95+
FormulaOrTerm::BinOp(lhs, binop, rhs) => {
96+
let pred = match binop {
97+
AmbiguousBinOp::Eq => chc::KnownPred::EQUAL,
98+
AmbiguousBinOp::Ne => chc::KnownPred::NOT_EQUAL,
99+
AmbiguousBinOp::Ge => chc::KnownPred::GREATER_THAN_OR_EQUAL,
100+
AmbiguousBinOp::Le => chc::KnownPred::LESS_THAN_OR_EQUAL,
101+
AmbiguousBinOp::Gt => chc::KnownPred::GREATER_THAN,
102+
AmbiguousBinOp::Lt => chc::KnownPred::LESS_THAN,
103+
};
104+
chc::Formula::Atom(chc::Atom::new(pred.into(), vec![lhs, rhs]))
105+
}
106+
FormulaOrTerm::Not(formula_or_term) => formula_or_term.into_formula()?.not(),
107+
FormulaOrTerm::Literal(b) => {
108+
if b {
109+
chc::Formula::top()
110+
} else {
111+
chc::Formula::bottom()
112+
}
113+
}
114+
};
115+
Some(fo)
116+
}
117+
118+
fn into_term(self) -> Option<chc::Term<T>> {
119+
let t = match self {
120+
FormulaOrTerm::Formula { .. } => return None,
121+
FormulaOrTerm::Term(t) => t,
122+
FormulaOrTerm::BinOp(lhs, AmbiguousBinOp::Eq, rhs) => lhs.eq(rhs),
123+
FormulaOrTerm::BinOp(lhs, AmbiguousBinOp::Ne, rhs) => lhs.ne(rhs),
124+
FormulaOrTerm::BinOp(lhs, AmbiguousBinOp::Ge, rhs) => lhs.ge(rhs),
125+
FormulaOrTerm::BinOp(lhs, AmbiguousBinOp::Le, rhs) => lhs.le(rhs),
126+
FormulaOrTerm::BinOp(lhs, AmbiguousBinOp::Gt, rhs) => lhs.gt(rhs),
127+
FormulaOrTerm::BinOp(lhs, AmbiguousBinOp::Lt, rhs) => lhs.lt(rhs),
128+
FormulaOrTerm::Not(formula_or_term) => formula_or_term.into_term()?.not(),
129+
FormulaOrTerm::Literal(b) => chc::Term::bool(b),
130+
};
131+
Some(t)
132+
}
133+
}
134+
135+
pub struct AnnotFnTranslator<'tcx> {
136+
tcx: TyCtxt<'tcx>,
137+
local_def_id: LocalDefId,
138+
139+
typeck: &'tcx mir_ty::TypeckResults<'tcx>,
140+
body: &'tcx rustc_hir::Body<'tcx>,
141+
142+
env: HashMap<HirId, chc::Term<rty::FunctionParamIdx>>,
143+
}
144+
145+
impl<'tcx> AnnotFnTranslator<'tcx> {
146+
pub fn new(tcx: TyCtxt<'tcx>, local_def_id: LocalDefId) -> Self {
147+
let map = tcx.hir();
148+
let body_id = map.body_owned_by(local_def_id);
149+
let body = map.body(body_id);
150+
151+
let typeck = tcx.typeck(local_def_id);
152+
let mut translator = Self {
153+
tcx,
154+
local_def_id,
155+
typeck,
156+
body,
157+
env: HashMap::default(),
158+
};
159+
translator.build_env_from_params();
160+
translator
161+
}
162+
163+
fn build_env_from_params(&mut self) {
164+
for (idx, param) in self.body.params.iter().enumerate() {
165+
let param_idx = rty::FunctionParamIdx::from(idx);
166+
let term = chc::Term::var(param_idx);
167+
self.build_env_from_pat(term, &param.pat);
168+
}
169+
}
170+
171+
fn build_env_from_pat(
172+
&mut self,
173+
param: chc::Term<rty::FunctionParamIdx>,
174+
pat: &'tcx rustc_hir::Pat<'tcx>,
175+
) {
176+
use rustc_hir::PatKind;
177+
178+
match pat.kind {
179+
PatKind::Binding(_, hir_id, _, None) => {
180+
self.env.insert(hir_id, param);
181+
}
182+
PatKind::TupleStruct(_, subpats, _) | PatKind::Tuple(subpats, _) => {
183+
for (idx, subpat) in subpats.iter().enumerate() {
184+
let field_term = param.clone().tuple_proj(idx.into());
185+
self.build_env_from_pat(field_term, subpat);
186+
}
187+
}
188+
_ => unimplemented!("unsupported pattern in formula: {:?}", pat),
189+
}
190+
}
191+
192+
pub fn to_formula_fn(&self) -> FormulaFn<'tcx> {
193+
let formula = self.to_formula(self.body.value);
194+
let params = self
195+
.tcx
196+
.fn_sig(self.local_def_id.to_def_id())
197+
.instantiate_identity()
198+
.skip_binder()
199+
.inputs()
200+
.to_vec();
201+
FormulaFn {
202+
params: IndexVec::from_raw(params),
203+
formula,
204+
}
205+
}
206+
207+
fn to_formula(&self, hir: &'tcx rustc_hir::Expr<'tcx>) -> chc::Formula<rty::FunctionParamIdx> {
208+
self.to_formula_or_term(hir)
209+
.into_formula()
210+
.expect("expected a formula")
211+
}
212+
213+
fn to_term(&self, hir: &'tcx rustc_hir::Expr<'tcx>) -> chc::Term<rty::FunctionParamIdx> {
214+
self.to_formula_or_term(hir)
215+
.into_term()
216+
.expect("expected a term")
217+
}
218+
219+
fn to_formula_or_term(
220+
&self,
221+
hir: &'tcx rustc_hir::Expr<'tcx>,
222+
) -> FormulaOrTerm<rty::FunctionParamIdx> {
223+
use rustc_hir::ExprKind;
224+
225+
match hir.kind {
226+
ExprKind::Binary(op, lhs, rhs) => {
227+
match op.node {
228+
rustc_hir::BinOpKind::Or => {
229+
let lhs = self.to_formula(lhs);
230+
let rhs = self.to_formula(rhs);
231+
return FormulaOrTerm::Formula(lhs.or(rhs));
232+
}
233+
rustc_hir::BinOpKind::And => {
234+
let lhs = self.to_formula(lhs);
235+
let rhs = self.to_formula(rhs);
236+
return FormulaOrTerm::Formula(lhs.and(rhs));
237+
}
238+
_ => {}
239+
}
240+
241+
let binop = match op.node {
242+
rustc_hir::BinOpKind::Eq => AmbiguousBinOp::Eq,
243+
rustc_hir::BinOpKind::Ne => AmbiguousBinOp::Ne,
244+
rustc_hir::BinOpKind::Ge => AmbiguousBinOp::Ge,
245+
rustc_hir::BinOpKind::Le => AmbiguousBinOp::Le,
246+
rustc_hir::BinOpKind::Gt => AmbiguousBinOp::Gt,
247+
rustc_hir::BinOpKind::Lt => AmbiguousBinOp::Lt,
248+
_ => unimplemented!("unsupported binary operator in formula: {:?}", op),
249+
};
250+
let lhs = self.to_formula_or_term(lhs);
251+
let rhs = self.to_formula_or_term(rhs);
252+
FormulaOrTerm::BinOp(lhs.into_term().unwrap(), binop, rhs.into_term().unwrap())
253+
}
254+
ExprKind::Unary(op, operand) => match op {
255+
rustc_hir::UnOp::Neg => {
256+
let operand = self.to_term(operand);
257+
FormulaOrTerm::Term(operand.neg())
258+
}
259+
rustc_hir::UnOp::Not => {
260+
FormulaOrTerm::Not(Box::new(self.to_formula_or_term(operand)))
261+
}
262+
_ => unimplemented!("unsupported unary operator in formula: {:?}", op),
263+
},
264+
ExprKind::Lit(lit) => match lit.node {
265+
rustc_ast::LitKind::Int(i, _) => {
266+
let t = chc::Term::int(i.get() as i64);
267+
FormulaOrTerm::Term(t)
268+
}
269+
rustc_ast::LitKind::Bool(b) => FormulaOrTerm::Literal(b),
270+
_ => unimplemented!("unsupported literal in formula: {:?}", lit),
271+
},
272+
ExprKind::Path(qpath) => {
273+
if let rustc_hir::def::Res::Local(hir_id) =
274+
self.typeck.qpath_res(&qpath, self.body.value.hir_id)
275+
{
276+
FormulaOrTerm::Term(
277+
self.env
278+
.get(&hir_id)
279+
.expect("unbound variable in formula")
280+
.clone(),
281+
)
282+
} else {
283+
unimplemented!("unsupported path in formula: {:?}", qpath);
284+
}
285+
}
286+
ExprKind::Block(block, _) => {
287+
if block.stmts.is_empty() {
288+
self.to_formula_or_term(block.expr.expect("expected an expression in block"))
289+
} else {
290+
unimplemented!("unsupported block in formula: {:?}", block);
291+
}
292+
}
293+
_ => unimplemented!("unsupported expression in formula: {:?}", hir),
294+
}
295+
}
296+
}

0 commit comments

Comments
 (0)