charon_driver/translate/
translate_closures.rs

1//! In rust, closures behave like ADTs that implement the FnOnce/FnMut/Fn traits automatically.
2//! Function item types also implement those traits automatically.
3//!
4//! Here we convert closures to a struct containing the closure's state (upvars), along with
5//! matching trait impls and fun decls (e.g. a Fn closure will have a trait impl for Fn, FnMut and
6//! FnOnce, along with 3 matching method implementations for call, call_mut and call_once).
7//! Function item types reuse the same generated trait impls and methods, with an empty state that
8//! forwards to the original function item.
9//!
10//! For example, given the following Rust code:
11//! ```ignore
12//! pub fn test_closure_capture<T: Clone>() {
13//!     let mut v = vec![];
14//!     let mut add = |x: &u32| v.push(*x);
15//!     add(&0);
16//!     add(&1);
17//! }
18//! ```
19//!
20//! We generate the equivalent desugared code:
21//! ```text
22//! struct {test_closure_capture::closure#0}<'a, T: Clone> (&'a mut Vec<u32>);
23//!
24//! // The 'a comes from captured variables, the 'b comes from the closure higher-kinded signature.
25//! impl<'a, 'b, T: Clone> FnMut<(&'b u32,)> for {test_closure_capture::closure#0}<'a, T> {
26//!     fn call_mut<'c>(&'c mut self, arg: (&'b u32,)) {
27//!         self.0.push(*arg.0);
28//!     }
29//! }
30//!
31//! impl<'a, 'b, T: Clone> FnOnce<(&'b u32,)> for {test_closure_capture::closure#0}<'a, T> {
32//!     type Output = ();
33//!     ...
34//! }
35//!
36//! pub fn test_closure_capture<T: Clone>() {
37//!     let mut v = vec![];
38//!     let mut add = {test_closure_capture::closure#0} (&mut v);
39//!     state.call_mut(&0);
40//!     state.call_mut(&1);
41//! }
42//! ```
43
44use crate::hax;
45use itertools::Itertools;
46use std::mem;
47
48use super::translate_crate::TransItemSourceKind;
49use super::translate_ctx::*;
50use charon_lib::ullbc_ast::*;
51
52pub fn translate_closure_kind(kind: &hax::ClosureKind) -> ClosureKind {
53    match kind {
54        hax::ClosureKind::Fn => ClosureKind::Fn,
55        hax::ClosureKind::FnMut => ClosureKind::FnMut,
56        hax::ClosureKind::FnOnce => ClosureKind::FnOnce,
57    }
58}
59
60/// If this trait proof is a built-in impl of a `Fn*` trait, return the `Self` type it is
61/// implemented for and the kind of the implemented trait.
62pub fn recognize_fn_trait_impl_proof(
63    trait_proof: &hax::TraitProof,
64) -> Option<(hax::Binder<&hax::Ty>, ClosureKind)> {
65    let hax::TraitProofKind::Builtin {
66        trait_data: hax::BuiltinTraitData::Other(lang_item),
67        ..
68    } = &trait_proof.kind
69    else {
70        return None;
71    };
72    let kind = match lang_item {
73        hax::SolverTraitLangItem::FnOnce => ClosureKind::FnOnce,
74        hax::SolverTraitLangItem::FnMut => ClosureKind::FnMut,
75        hax::SolverTraitLangItem::Fn => ClosureKind::Fn,
76        _ => return None,
77    };
78    let Some(hax::GenericArg::Type(self_ty)) =
79        trait_proof.pred.hax_skip_binder_ref().generic_args.first()
80    else {
81        unreachable!("no `Self` type arg on a `Fn*` trait ref")
82    };
83    Some((trait_proof.pred.rebind(self_ty), kind))
84}
85
86/// The built-in `Fn*` impl of the given kind that we generate for this closure or function item.
87pub fn callable_virtual_impl<'a>(
88    def: &'a hax::FullDef<'_>,
89    target_kind: ClosureKind,
90) -> &'a hax::VirtualTraitImpl {
91    CallableFnImpls::from_def(def)
92        .and_then(|impls| impls.vimpl(target_kind))
93        .expect("expected a callable with a Fn* impl")
94}
95
96#[derive(Clone, Copy)]
97enum Callable<'a> {
98    Closure(&'a hax::ClosureArgs),
99    FnDef {
100        item: &'a hax::ItemRef,
101        sig: &'a hax::PolyFnSig,
102        /// The arguments, tupled as the `Fn*` traits take them. Binds the same variables as `sig`.
103        tupled_args_ty: &'a hax::Binder<hax::Ty>,
104    },
105}
106
107impl<'a> Callable<'a> {
108    fn item(self) -> &'a hax::ItemRef {
109        match self {
110            Callable::Closure(args) => &args.item,
111            Callable::FnDef { item, .. } => item,
112        }
113    }
114
115    fn sig(self) -> &'a hax::PolyFnSig {
116        match self {
117            Callable::Closure(args) => &args.fn_sig,
118            Callable::FnDef { sig, .. } => sig,
119        }
120    }
121
122    /// The arguments, tupled as the `Fn*` traits take them, e.g. `(A, B, C)`. This is under the
123    /// same binder as `sig`.
124    fn tupled_args_ty(self) -> &'a hax::Ty {
125        match self {
126            Callable::Closure(args) => args.tupled_args_ty.hax_skip_binder_ref(),
127            Callable::FnDef { tupled_args_ty, .. } => tupled_args_ty.hax_skip_binder_ref(),
128        }
129    }
130}
131
132#[derive(Clone, Copy)]
133struct CallableFnImpls<'a> {
134    callable: Callable<'a>,
135    fn_once_impl: Option<&'a hax::VirtualTraitImpl>,
136    fn_mut_impl: Option<&'a hax::VirtualTraitImpl>,
137    fn_impl: Option<&'a hax::VirtualTraitImpl>,
138}
139
140impl<'a> CallableFnImpls<'a> {
141    fn from_def(def: &'a hax::FullDef<'_>) -> Option<Self> {
142        match def.kind() {
143            hax::FullDefKind::Closure {
144                args,
145                fn_once_impl,
146                fn_mut_impl,
147                fn_impl,
148                ..
149            } => Some(Self {
150                callable: Callable::Closure(args),
151                fn_once_impl: Some(fn_once_impl),
152                fn_mut_impl: fn_mut_impl.as_deref(),
153                fn_impl: fn_impl.as_deref(),
154            }),
155            hax::FullDefKind::Fn {
156                sig,
157                tupled_args_ty,
158                fn_once_impl,
159                fn_mut_impl,
160                fn_impl,
161                ..
162            }
163            | hax::FullDefKind::AssocFn {
164                sig,
165                tupled_args_ty,
166                fn_once_impl,
167                fn_mut_impl,
168                fn_impl,
169                ..
170            }
171            | hax::FullDefKind::Ctor {
172                sig,
173                tupled_args_ty,
174                fn_once_impl,
175                fn_mut_impl,
176                fn_impl,
177                ..
178            } => Some(Self {
179                callable: Callable::FnDef {
180                    item: def.this(),
181                    sig,
182                    tupled_args_ty: tupled_args_ty.as_ref()?,
183                },
184                fn_once_impl: fn_once_impl.as_deref(),
185                fn_mut_impl: fn_mut_impl.as_deref(),
186                fn_impl: fn_impl.as_deref(),
187            }),
188            _ => None,
189        }
190    }
191
192    fn vimpl(self, target_kind: ClosureKind) -> Option<&'a hax::VirtualTraitImpl> {
193        match target_kind {
194            ClosureKind::FnOnce => self.fn_once_impl,
195            ClosureKind::FnMut => self.fn_mut_impl,
196            ClosureKind::Fn => self.fn_impl,
197        }
198    }
199}
200
201/// References to callable items are subtle because there are three sources of lifetimes on top of
202/// the normal generics: closure upvars, the higher-kindedness of the callable itself, and the
203/// late-bound generics of the `call`/`call_mut` methods. One must be careful to choose the right
204/// method from these.
205impl<'tcx> ItemTransCtx<'tcx, '_> {
206    /// Translate a reference to a callable item that takes late-bound lifetimes. The binder binds
207    /// the late-bound lifetimes of the callable itself, if it is higher-kinded.
208    fn translate_callable_bound_ref_with_late_bound(
209        &mut self,
210        span: Span,
211        callable: Callable<'_>,
212        kind: TransItemSourceKind,
213    ) -> Result<RegionBinder<DeclRef<ItemId>>, Error> {
214        if !matches!(
215            kind,
216            TransItemSourceKind::TraitImpl(..) | TransItemSourceKind::ClosureAsFnCast
217        ) {
218            raise_error!(
219                self,
220                span,
221                "Called `translate_callable_bound_ref_with_late_bound` on a `{kind:?}`; \
222                use `translate_closure_ref_with_upvars` \
223                or `translate_callable_bound_ref_with_method_bound` instead"
224            )
225        }
226        let dref: DeclRef<ItemId> = self.translate_item(span, callable.item(), kind)?;
227        self.translate_region_binder(span, callable.sig(), |ctx, _| {
228            let mut dref = dref.move_under_binder();
229            // The regions for these item kinds have the fn late bound regions at the end.
230            for (a, b) in dref.generics.regions.iter_mut().rev().zip(
231                ctx.innermost_binder()
232                    .params
233                    .identity_args()
234                    .regions
235                    .into_iter()
236                    .rev(),
237            ) {
238                *a = b;
239            }
240            Ok(dref)
241        })
242    }
243
244    /// Translate a reference to a callable item that takes late-bound lifetimes and method
245    /// lifetimes. The binder binds the late-bound lifetimes of the `call`/`call_mut` method
246    /// (specified by `target_kind`).
247    fn translate_callable_bound_ref_with_method_bound(
248        &mut self,
249        span: Span,
250        item: &hax::ItemRef,
251        kind: TransItemSourceKind,
252        target_kind: ClosureKind,
253    ) -> Result<RegionBinder<DeclRef<ItemId>>, Error> {
254        if !matches!(kind, TransItemSourceKind::CallableMethod(..)) {
255            raise_error!(
256                self,
257                span,
258                "Called `translate_callable_bound_ref_with_method_bound` on a `{kind:?}`; \
259                use `translate_closure_ref_with_upvars` \
260                or `translate_callable_bound_ref_with_late_bound` instead"
261            )
262        }
263        let dref: DeclRef<ItemId> = self.translate_item(span, item, kind)?;
264        let mut dref = dref.move_under_binder();
265        let mut regions = IndexVec::new();
266        match target_kind {
267            ClosureKind::FnOnce => {}
268            ClosureKind::FnMut | ClosureKind::Fn => {
269                let rid =
270                    regions.push_with(|index| RegionParam::new(index, None, Variance::Covariant));
271                *dref.generics.regions.iter_mut().last().unwrap() =
272                    Region::Var(DeBruijnVar::new_at_zero(rid));
273            }
274        }
275        Ok(RegionBinder {
276            regions,
277            skip_binder: dref,
278        })
279    }
280
281    /// If this trait proof is the built-in impl of a `Fn*` trait for a closure or function item,
282    /// return the callable item and the kind of the implemented trait. The returned item has its
283    /// regions erased.
284    pub fn recognize_callable_impl_proof(
285        &self,
286        trait_proof: &hax::TraitProof,
287    ) -> Option<(hax::ItemRef, ClosureKind)> {
288        let hax::TraitProofKind::Builtin {
289            trait_data: hax::BuiltinTraitData::Other(lang_item),
290            ..
291        } = &trait_proof.kind
292        else {
293            return None;
294        };
295        let target_kind = match lang_item {
296            hax::SolverTraitLangItem::FnOnce => ClosureKind::FnOnce,
297            hax::SolverTraitLangItem::FnMut => ClosureKind::FnMut,
298            hax::SolverTraitLangItem::Fn => ClosureKind::Fn,
299            _ => return None,
300        };
301        // We skip the binder and erase regions to avoid bound vars escaping.
302        let hax::GenericArg::Type(callable_ty) = trait_proof
303            .pred
304            .hax_skip_binder_ref()
305            .generic_args
306            .first()?
307        else {
308            return None;
309        };
310        let item = match callable_ty.kind() {
311            hax::TyKind::Closure(closure_args) => &closure_args.item,
312            hax::TyKind::FnDef { item, .. } => item,
313            _ => return None,
314        };
315        Some((item.erase(self.hax_state_with_id()), target_kind))
316    }
317
318    pub(crate) fn translate_callable_method_fn_ptr(
319        &mut self,
320        span: Span,
321        item: &hax::ItemRef,
322    ) -> Result<Option<RegionBinder<FnPtr>>, Error> {
323        if !self.monomorphize() {
324            return Ok(None);
325        }
326        let Some(in_trait) = &item.in_trait else {
327            return Ok(None);
328        };
329        let Some((callable, target_kind)) = self.recognize_callable_impl_proof(in_trait) else {
330            return Ok(None);
331        };
332        let kind = TransItemSourceKind::CallableMethod(target_kind);
333        let bound_ref = self.translate_callable_bound_ref_with_method_bound(
334            span,
335            &callable,
336            kind,
337            target_kind,
338        )?;
339        Ok(Some(bound_ref.map(|dref| {
340            let fn_ref: FunDeclRef = dref.try_into().unwrap();
341            FnPtr::new(FnPtrKind::Fun(fn_ref.id), fn_ref.generics)
342        })))
343    }
344}
345
346impl<'tcx> ItemTransCtx<'tcx, '_> {
347    /// Translate a reference to the closure ADT.
348    pub fn translate_closure_type_ref(
349        &mut self,
350        span: Span,
351        closure: &hax::ClosureArgs,
352    ) -> Result<TypeDeclRef, Error> {
353        self.translate_type_decl_ref(span, &closure.item)
354    }
355
356    /// For stateless closures, translate a function reference to the top-level function that
357    /// executes the closure code without taking the state as parameter.If you want to instantiate
358    /// the binder, use the lifetimes from `self.closure_late_regions`.
359    pub fn translate_stateless_closure_as_fn_ref(
360        &mut self,
361        span: Span,
362        closure: &hax::ClosureArgs,
363    ) -> Result<RegionBinder<FunDeclRef>, Error> {
364        let kind = TransItemSourceKind::ClosureAsFnCast;
365        let bound_dref = self.translate_callable_bound_ref_with_late_bound(
366            span,
367            Callable::Closure(closure),
368            kind,
369        )?;
370        Ok(bound_dref.map(|dref| dref.try_into().unwrap()))
371    }
372
373    /// Translate a reference to the chosen closure impl. The resulting value needs lifetime
374    /// arguments for late-bound lifetimes. If you want to instantiate the binder, use the
375    /// lifetimes from `self.closure_late_regions`.
376    pub fn translate_closure_bound_impl_ref(
377        &mut self,
378        span: Span,
379        closure: &hax::ClosureArgs,
380        target_kind: ClosureKind,
381    ) -> Result<RegionBinder<TraitImplRef>, Error> {
382        let kind = TransItemSourceKind::TraitImpl(TransImplSource::Callable(target_kind));
383        let bound_dref = self.translate_callable_bound_ref_with_late_bound(
384            span,
385            Callable::Closure(closure),
386            kind,
387        )?;
388        Ok(bound_dref.map(|dref| dref.try_into().unwrap()))
389    }
390
391    /// Translate a reference to the chosen callable impl.
392    pub fn translate_callable_impl_ref(
393        &mut self,
394        span: Span,
395        item: &hax::ItemRef,
396        target_kind: ClosureKind,
397    ) -> Result<TraitImplRef, Error> {
398        self.translate_item(
399            span,
400            item,
401            TransItemSourceKind::TraitImpl(TransImplSource::Callable(target_kind)),
402        )
403    }
404
405    pub fn translate_closure_info(
406        &mut self,
407        span: Span,
408        args: &hax::ClosureArgs,
409    ) -> Result<ClosureInfo, Error> {
410        use ClosureKind::*;
411        let kind = translate_closure_kind(&args.kind);
412
413        let fn_once_impl = self.translate_closure_bound_impl_ref(span, args, FnOnce)?;
414        let fn_mut_impl = if matches!(kind, FnMut | Fn) {
415            Some(self.translate_closure_bound_impl_ref(span, args, FnMut)?)
416        } else {
417            None
418        };
419        let fn_impl = if matches!(kind, Fn) {
420            Some(self.translate_closure_bound_impl_ref(span, args, Fn)?)
421        } else {
422            None
423        };
424        let signature = self.translate_poly_fun_sig(span, &args.fn_sig)?;
425        Ok(ClosureInfo {
426            kind,
427            fn_once_impl,
428            fn_mut_impl,
429            fn_impl,
430            signature,
431        })
432    }
433
434    fn get_callable_state_ty(&mut self, span: Span, callable: Callable<'_>) -> Result<Ty, Error> {
435        Ok(match callable {
436            Callable::Closure(args) => {
437                let tref = self.translate_closure_type_ref(span, args)?;
438                TyKind::Adt(tref).into_ty()
439            }
440            Callable::FnDef { item, .. } => {
441                let fn_ref = self.translate_bound_fn_ptr(span, item, TransItemSourceKind::Fun)?;
442                TyKind::FnDef(fn_ref).into_ty()
443            }
444        })
445    }
446
447    /// Translate the types of the captured variables. Should be called only in
448    /// `translate_item_generics`. If you need these types, fetch them in
449    /// `outermost_binder().closure_upvar_tys`.
450    pub fn translate_closure_upvar_tys(
451        &mut self,
452        span: Span,
453        args: &hax::ClosureArgs,
454    ) -> Result<IndexVec<FieldId, Ty>, Error> {
455        args.upvar_tys
456            .iter()
457            .map(|ty| self.translate_ty(span, ty))
458            .try_collect()
459    }
460
461    pub fn translate_closure_adt(
462        &mut self,
463        span: Span,
464        _args: &hax::ClosureArgs,
465    ) -> Result<TypeDeclKind, Error> {
466        let fields: IndexVec<FieldId, Field> = self
467            .the_only_binder()
468            .closure_upvar_tys
469            .as_ref()
470            .unwrap()
471            .iter()
472            .cloned()
473            .enumerate()
474            .map(|(field_id, ty)| Field {
475                span,
476                attr_info: AttrInfo::dummy_private(),
477                name: format!("_{field_id}"),
478                is_positional: true,
479                ty,
480            })
481            .collect();
482        Ok(TypeDeclKind::Struct(fields))
483    }
484
485    /// Given an item that is callable, generate the signature of the
486    /// `call_once`/`call_mut`/`call` method (depending on `target_kind`).
487    fn translate_callable_method_sig(
488        &mut self,
489        def: &hax::FullDef<'tcx>,
490        span: Span,
491        callable: Callable,
492        target_kind: ClosureKind,
493    ) -> Result<RegionBinder<FunSig>, Error> {
494        let signature = callable.sig();
495        trace!(
496            "signature of callable {:?}:\n{:?}",
497            def.def_id(),
498            signature.value,
499        );
500
501        let mut bound_regions = IndexVec::new();
502        let mut fun_sig = self
503            .translate_fun_sig(span, signature.hax_skip_binder_ref())?
504            .move_under_binder();
505        let state_ty = self
506            .get_callable_state_ty(span, callable)?
507            .move_under_binder();
508
509        // Depending on the kind of the closure generated, add a reference
510        let state_ty = match target_kind {
511            ClosureKind::FnOnce => state_ty,
512            ClosureKind::Fn | ClosureKind::FnMut => {
513                let rid = bound_regions
514                    .push_with(|index| RegionParam::new(index, None, Variance::Covariant));
515                let r = Region::Var(DeBruijnVar::new_at_zero(rid));
516                let mutability = if target_kind == ClosureKind::Fn {
517                    RefKind::Shared
518                } else {
519                    RefKind::Mut
520                };
521                TyKind::Ref(r, state_ty, mutability).into_ty()
522            }
523        };
524
525        let tupled_args_ty = self
526            .translate_ty(span, callable.tupled_args_ty())?
527            .move_under_binder();
528        fun_sig.inputs = vec![state_ty, tupled_args_ty];
529
530        Ok(RegionBinder {
531            regions: bound_regions,
532            skip_binder: fun_sig,
533        })
534    }
535
536    fn translate_callable_method_body(
537        &mut self,
538        span: Span,
539        def: &hax::FullDef<'tcx>,
540        target_kind: ClosureKind,
541        callable: Callable,
542        signature: &FunSig,
543    ) -> Result<Body, Error> {
544        match callable {
545            Callable::Closure(args) => {
546                self.translate_closure_method_body(span, def, target_kind, args, signature)
547            }
548            Callable::FnDef { item, .. } => {
549                self.translate_fn_def_method_body(span, item, signature)
550            }
551        }
552    }
553
554    fn translate_closure_method_body(
555        &mut self,
556        span: Span,
557        def: &hax::FullDef<'tcx>,
558        target_kind: ClosureKind,
559        args: &hax::ClosureArgs,
560        signature: &FunSig,
561    ) -> Result<Body, Error> {
562        use ClosureKind::*;
563        let closure_kind = translate_closure_kind(&args.kind);
564        Ok(match (target_kind, closure_kind) {
565            (Fn, Fn) | (FnMut, FnMut) | (FnOnce, FnOnce) => {
566                // Translate the function's body normally
567                let mut body = self.translate_def_body(span, def);
568                // The body is translated as if the locals are: ret value, state, arg-1,
569                // ..., arg-N, rest...
570                // However, there is only one argument with the tupled closure arguments;
571                // we must thus shift all locals with index >=2 by 1, and add a new local
572                // for the tupled arg, giving us: ret value, state, args, arg-1, ...,
573                // arg-N, rest...
574                // We then add N statements of the form `locals[N+3] := move locals[2].N`,
575                // to destructure the arguments.
576                let Body::Unstructured(GExprBody {
577                    locals,
578                    body: blocks,
579                    ..
580                }) = &mut body
581                else {
582                    return Ok(body);
583                };
584
585                // The (Arg1, Arg2, ..) type.
586                let tupled_ty = &signature.inputs[1];
587
588                blocks.dyn_visit_mut(|local: &mut LocalId| {
589                    if local.index() >= 2 {
590                        *local += 1;
591                    }
592                });
593
594                // Remember how many arguments there are
595                let closure_arg_count = locals.arg_count - 1;
596                let mut old_locals = mem::take(&mut locals.locals).into_iter();
597                locals.arg_count = 2;
598                locals.locals.push(old_locals.next().unwrap()); // ret
599                locals.locals.push(old_locals.next().unwrap()); // state
600                let tupled_arg = locals.new_var(Some("tupled_args".to_string()), tupled_ty.clone());
601                locals.locals.extend(old_locals.map(|mut l| {
602                    l.index += 1;
603                    l
604                }));
605
606                let untupled_args = locals
607                    .locals
608                    .iter()
609                    .skip(3)
610                    .take(closure_arg_count)
611                    .map(|l| &l.ty)
612                    .cloned();
613                let new_stts = untupled_args.enumerate().map(|(i, ty)| {
614                    let nth_field = tupled_arg
615                        .clone()
616                        .project(ProjectionElem::Field(None, FieldId::new(i)), ty);
617                    let local_id = LocalId::new(i + 3);
618                    Statement::new(
619                        span,
620                        StatementKind::Assign(
621                            locals.place_for_var(local_id),
622                            Rvalue::Use(Operand::Move(nth_field), WithRetag::No),
623                        ),
624                    )
625                });
626                blocks[BlockId::ZERO].statements.splice(0..0, new_stts);
627
628                body
629            }
630            // Target translation:
631            //
632            // fn call_once(state: Self, args: Args) -> Output {
633            //   let temp_ref = &[mut] state;
634            //   let ret = self.call[_mut](temp, args);
635            //   drop state;
636            //   return ret;
637            // }
638            //
639            (FnOnce, Fn | FnMut) => {
640                // Hax (via rustc) gives us the MIR to do this.
641                let Some(body) = def.this.closure_once_shim(self.hax_state()) else {
642                    panic!("missing shim for closure")
643                };
644                self.translate_body(span, body, &def.source_text)
645            }
646            // Target translation:
647            //
648            // fn call_mut(state: &mut Self, args: Args) -> Output {
649            //   let reborrow = &*state;
650            //   self.call(reborrow, args)
651            // }
652            (FnMut, Fn) => {
653                let fun_id: FunDeclId = self.register_item(
654                    span,
655                    def.this(),
656                    TransItemSourceKind::CallableMethod(closure_kind),
657                );
658                let impl_ref = self.translate_callable_impl_ref(span, &args.item, closure_kind)?;
659                // TODO: make a trait call to avoid needing to concatenate things ourselves.
660                // TODO: can we ask hax for the trait ref?
661                let fn_op = FnOperand::Regular(FnPtr::new(
662                    fun_id.into(),
663                    impl_ref.generics.concat(&GenericArgs {
664                        regions: vec![self.translate_erased_region()].into(),
665                        ..GenericArgs::empty()
666                    }),
667                ));
668
669                let mut builder = BodyBuilder::new(span, 2);
670
671                let output = builder.new_var(None, signature.output.clone());
672                let state = builder.new_var(Some("state".to_string()), signature.inputs[0].clone());
673                let args = builder.new_var(Some("args".to_string()), signature.inputs[1].clone());
674                let deref_state = state.deref();
675                let reborrow_ty = TyKind::Ref(
676                    self.translate_erased_region(),
677                    deref_state.ty.clone(),
678                    RefKind::Shared,
679                )
680                .into_ty();
681                let reborrow = builder.new_var(None, reborrow_ty);
682
683                builder.push_statement(StatementKind::Assign(
684                    reborrow.clone(),
685                    Rvalue::Ref {
686                        place: deref_state,
687                        kind: BorrowKind::Shared,
688                        // The state must be Sized, hence `()` as ptr-metadata
689                        ptr_metadata: Operand::mk_const_unit(),
690                    },
691                ));
692
693                builder.call(Call {
694                    func: fn_op,
695                    args: vec![Operand::Move(reborrow), Operand::Move(args)],
696                    dest: output,
697                });
698
699                Body::Unstructured(builder.build())
700            }
701            (Fn, FnOnce) | (Fn, FnMut) | (FnMut, FnOnce) => {
702                panic!(
703                    "Can't make a closure body for a more restrictive kind \
704                    than the closure kind"
705                )
706            }
707        })
708    }
709
710    fn translate_fn_def_method_body(
711        &mut self,
712        span: Span,
713        item: &hax::ItemRef,
714        signature: &FunSig,
715    ) -> Result<Body, Error> {
716        let late_bound_regions = self
717            .innermost_binder()
718            .bound_region_vars
719            .iter()
720            .map(|rid| Region::Var(DeBruijnVar::new_at_zero(*rid)))
721            .collect();
722        let fn_ptr = self
723            .translate_bound_fn_ptr(span, item, TransItemSourceKind::Fun)?
724            .apply(late_bound_regions);
725        let fn_op = FnOperand::Regular(fn_ptr);
726
727        let mut builder = BodyBuilder::new(span, 2);
728
729        let output = builder.new_var(None, signature.output.clone());
730        let _state = builder.new_var(Some("state".to_string()), signature.inputs[0].clone());
731        let tupled_args_ty = &signature.inputs[1];
732        let tupled_args = builder.new_var(Some("args".to_string()), tupled_args_ty.clone());
733
734        // We need the type declaration to have been translated to get the fields (since in monomorphic)
735        // mode they aren't in the generics. So ensure it's been translated!
736        let tuple_id = tupled_args_ty.as_adt().unwrap().id;
737        let _ = self.get_or_translate(ItemId::Type(tuple_id))?;
738        let arg_tys = tupled_args_ty.as_tuple_fields(&self.t_ctx.translated);
739
740        let args = arg_tys
741            .into_iter()
742            .enumerate()
743            .map(|(i, ty)| {
744                let nth_field = tupled_args
745                    .clone()
746                    .project(ProjectionElem::Field(None, FieldId::new(i)), ty);
747                Operand::Move(nth_field)
748            })
749            .collect();
750
751        builder.call(Call {
752            func: fn_op,
753            args,
754            dest: output,
755        });
756
757        Ok(Body::Unstructured(builder.build()))
758    }
759
760    /// Given an item that is a closure, generate the `call_once`/`call_mut`/`call` method
761    /// (depending on `target_kind`).
762    #[tracing::instrument(skip(self, item_meta))]
763    pub fn translate_closure_method(
764        mut self,
765        def_id: FunDeclId,
766        item_meta: ItemMeta,
767        def: &hax::FullDef<'tcx>,
768        target_kind: ClosureKind,
769    ) -> Result<FunDecl, Error> {
770        let span = item_meta.span;
771        let callable_impls = CallableFnImpls::from_def(def).unwrap();
772        let callable = callable_impls.callable;
773
774        // Hax gives us trait-related information for the impl we're building.
775        let vimpl = callable_impls.vimpl(target_kind).unwrap();
776        let implemented_trait = self.translate_trait_predicate(span, &vimpl.trait_pred)?;
777        let method_id =
778            self.translate_trait_method_id(implemented_trait.id, &vimpl.methods[0].0)?;
779
780        let impl_ref = self.translate_callable_impl_ref(span, callable.item(), target_kind)?;
781        let src = FunSource::TraitImpl {
782            impl_ref,
783            trait_ref: implemented_trait.clone(),
784            item_id: method_id,
785            reuses_default: false,
786        };
787
788        // Translate the function signature
789        let bound_sig = self.translate_callable_method_sig(def, span, callable, target_kind)?;
790        // We give it the lifetime parameter we had prepared for that purpose.
791        let signature = bound_sig.apply(
792            self.the_only_binder()
793                .closure_call_method_region
794                .iter()
795                .map(|r| Region::Var(DeBruijnVar::new_at_zero(*r)))
796                .collect(),
797        );
798
799        let body = if item_meta.opacity.with_private_contents().is_opaque() {
800            Body::Opaque
801        } else {
802            self.translate_callable_method_body(span, def, target_kind, callable, &signature)?
803        };
804
805        Ok(FunDecl {
806            def_id,
807            item_meta,
808            generics: self.into_generics(),
809            signature: Box::new(signature),
810            src,
811            body,
812        })
813    }
814
815    #[tracing::instrument(skip(self, item_meta))]
816    pub fn translate_closure_trait_impl(
817        mut self,
818        def_id: TraitImplId,
819        item_meta: ItemMeta,
820        def: &hax::FullDef<'tcx>,
821        target_kind: ClosureKind,
822    ) -> Result<TraitImpl, Error> {
823        let span = item_meta.span;
824        let callable_impls = CallableFnImpls::from_def(def).unwrap();
825        let callable = callable_impls.callable;
826
827        // Hax gives us trait-related information for the impl we're building.
828        let vimpl = callable_impls.vimpl(target_kind).unwrap();
829        let mut timpl = self.translate_virtual_trait_impl(
830            def_id,
831            item_meta,
832            callable.item(),
833            TransImplSource::Callable(target_kind),
834            vimpl,
835        )?;
836
837        // Construct the `call_*` method reference.
838        let trait_decl_id = timpl.impl_trait.id;
839        let trait_method_id = self.translate_trait_method_id(trait_decl_id, &vimpl.methods[0].0)?;
840        let call_fn_binder = {
841            let kind = TransItemSourceKind::CallableMethod(target_kind);
842            let bound_method_ref: RegionBinder<DeclRef<ItemId>> = self
843                .translate_callable_bound_ref_with_method_bound(
844                    span,
845                    callable.item(),
846                    kind,
847                    target_kind,
848                )?;
849            let params = GenericParams {
850                regions: bound_method_ref.regions,
851                ..GenericParams::empty()
852            };
853            let fn_decl_ref: FunDeclRef = bound_method_ref.skip_binder.try_into().unwrap();
854            Binder::new(
855                BinderKind::TraitMethod(trait_decl_id, trait_method_id),
856                params,
857                fn_decl_ref,
858            )
859        };
860        if self.monomorphize() {
861            return Ok(timpl);
862        }
863        timpl
864            .methods
865            .set_slot_extend(trait_method_id, call_fn_binder);
866
867        Ok(timpl)
868    }
869
870    /// Given an item that is a non-capturing closure, generate the equivalent function,
871    /// by removing the state from the parameters and untupling the arguments.
872    #[tracing::instrument(skip(self, item_meta))]
873    pub fn translate_stateless_closure_as_fn(
874        mut self,
875        def_id: FunDeclId,
876        item_meta: ItemMeta,
877        def: &hax::FullDef<'tcx>,
878    ) -> Result<FunDecl, Error> {
879        let span = item_meta.span;
880        let hax::FullDefKind::Closure { args: closure, .. } = &def.kind else {
881            unreachable!()
882        };
883
884        trace!("About to translate closure as fn:\n{:?}", def.def_id());
885
886        assert!(
887            closure.upvar_tys.is_empty(),
888            "Only stateless closures can be translated as functions"
889        );
890
891        // Translate the function signature
892        let signature = self.translate_fun_sig(span, closure.fn_sig.hax_skip_binder_ref())?;
893        let state_ty = self.get_callable_state_ty(span, Callable::Closure(closure))?;
894
895        let body = if item_meta.opacity.with_private_contents().is_opaque() {
896            Body::Opaque
897        } else {
898            // Target translation:
899            //
900            // fn call_fn(arg0: Args[0], ..., argN: Args[N]) -> Output {
901            //   let closure: Closure = {};
902            //   let args = (arg0, ..., argN);
903            //   closure.call(args)
904            // }
905            let fun_id: FunDeclId = self.register_item(
906                span,
907                def.this(),
908                TransItemSourceKind::CallableMethod(ClosureKind::FnOnce),
909            );
910            let impl_ref =
911                self.translate_callable_impl_ref(span, &closure.item, ClosureKind::FnOnce)?;
912            let fn_op = FnOperand::Regular(FnPtr::new(fun_id.into(), impl_ref.generics.clone()));
913
914            let mut builder = BodyBuilder::new(span, signature.inputs.len());
915
916            let output = builder.new_var(None, signature.output.clone());
917            let args: Vec<Place> = signature
918                .inputs
919                .iter()
920                .enumerate()
921                .map(|(i, ty)| builder.new_var(Some(format!("arg{}", i + 1)), ty.clone()))
922                .collect();
923            let args_tupled_ty =
924                self.translate_ty(span, closure.tupled_args_ty.hax_skip_binder_ref())?;
925            let args_tupled = builder.new_var(Some("args".to_string()), args_tupled_ty.clone());
926            let state = builder.new_var(Some("state".to_string()), state_ty.clone());
927
928            builder.push_statement(StatementKind::Assign(
929                args_tupled.clone(),
930                Rvalue::Aggregate(
931                    AggregateKind::Adt(args_tupled_ty.as_adt().unwrap().clone(), None, None),
932                    args.into_iter().map(Operand::Move).collect(),
933                ),
934            ));
935
936            let state_ty_adt = state_ty.as_adt().unwrap();
937            builder.push_statement(StatementKind::Assign(
938                state.clone(),
939                Rvalue::Aggregate(AggregateKind::Adt(state_ty_adt.clone(), None, None), vec![]),
940            ));
941
942            builder.call(Call {
943                func: fn_op,
944                args: vec![Operand::Move(state), Operand::Move(args_tupled)],
945                dest: output,
946            });
947
948            Body::Unstructured(builder.build())
949        };
950
951        Ok(FunDecl {
952            def_id,
953            item_meta,
954            generics: self.into_generics(),
955            signature: Box::new(signature),
956            src: FunSource::Normal,
957            body,
958        })
959    }
960}
charon_driver/translate/translate_closures.rs

charon_driver/translate/
translate_closures.rs
