1use crate::hax;
42use itertools::Itertools;
43use std::mem;
44
45use super::translate_crate::TransItemSourceKind;
46use super::translate_ctx::*;
47use charon_lib::ast::ullbc_ast_utils::BodyBuilder;
48use charon_lib::ast::*;
49use charon_lib::ids::IndexVec;
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
60impl<'tcx> ItemTransCtx<'tcx, '_> {
65 fn translate_closure_bound_ref_with_late_bound(
68 &mut self,
69 span: Span,
70 closure: &hax::ClosureArgs,
71 kind: TransItemSourceKind,
72 ) -> Result<RegionBinder<DeclRef<ItemId>>, Error> {
73 if !matches!(
74 kind,
75 TransItemSourceKind::TraitImpl(..) | TransItemSourceKind::ClosureAsFnCast
76 ) {
77 raise_error!(
78 self,
79 span,
80 "Called `translate_closure_bound_ref_with_late_bound` on a `{kind:?}`; \
81 use `translate_closure_ref_with_upvars` \
82 or `translate_closure_bound_ref_with_method_bound` instead"
83 )
84 }
85 let dref: DeclRef<ItemId> = self.translate_item(span, &closure.item, kind)?;
86 self.translate_region_binder(span, &closure.fn_sig, |ctx, _| {
87 let mut dref = dref.move_under_binder();
88 for (a, b) in dref.generics.regions.iter_mut().rev().zip(
90 ctx.innermost_binder()
91 .params
92 .identity_args()
93 .regions
94 .into_iter()
95 .rev(),
96 ) {
97 *a = b;
98 }
99 Ok(dref)
100 })
101 }
102
103 fn translate_closure_bound_ref_with_method_bound(
107 &mut self,
108 span: Span,
109 closure: &hax::ClosureArgs,
110 kind: TransItemSourceKind,
111 target_kind: ClosureKind,
112 ) -> Result<RegionBinder<DeclRef<ItemId>>, Error> {
113 if !matches!(kind, TransItemSourceKind::ClosureMethod(..)) {
114 raise_error!(
115 self,
116 span,
117 "Called `translate_closure_bound_ref_with_method_bound` on a `{kind:?}`; \
118 use `translate_closure_ref_with_upvars` \
119 or `translate_closure_bound_ref_with_late_bound` instead"
120 )
121 }
122 let dref: DeclRef<ItemId> = self.translate_item(span, &closure.item, kind)?;
123 let mut dref = dref.move_under_binder();
124 let mut regions = IndexVec::new();
125 match target_kind {
126 ClosureKind::FnOnce => {}
127 ClosureKind::FnMut | ClosureKind::Fn => {
128 let rid = regions.push_with(|index| RegionParam::new(index, None));
129 *dref.generics.regions.iter_mut().last().unwrap() =
130 Region::Var(DeBruijnVar::new_at_zero(rid));
131 }
132 }
133 Ok(RegionBinder {
134 regions,
135 skip_binder: dref,
136 })
137 }
138}
139
140impl<'tcx> ItemTransCtx<'tcx, '_> {
141 pub fn translate_closure_type_ref(
143 &mut self,
144 span: Span,
145 closure: &hax::ClosureArgs,
146 ) -> Result<TypeDeclRef, Error> {
147 self.translate_item(span, &closure.item, TransItemSourceKind::Type)
148 }
149
150 pub fn translate_stateless_closure_as_fn_ref(
154 &mut self,
155 span: Span,
156 closure: &hax::ClosureArgs,
157 ) -> Result<RegionBinder<FunDeclRef>, Error> {
158 let kind = TransItemSourceKind::ClosureAsFnCast;
159 let bound_dref = self.translate_closure_bound_ref_with_late_bound(span, closure, kind)?;
160 Ok(bound_dref.map(|dref| dref.try_into().unwrap()))
161 }
162
163 pub fn translate_closure_bound_impl_ref(
167 &mut self,
168 span: Span,
169 closure: &hax::ClosureArgs,
170 target_kind: ClosureKind,
171 ) -> Result<RegionBinder<TraitImplRef>, Error> {
172 let kind = TransItemSourceKind::TraitImpl(TraitImplSource::Closure(target_kind));
173 let bound_dref = self.translate_closure_bound_ref_with_late_bound(span, closure, kind)?;
174 Ok(bound_dref.map(|dref| dref.try_into().unwrap()))
175 }
176
177 pub fn translate_closure_impl_ref(
179 &mut self,
180 span: Span,
181 closure: &hax::ClosureArgs,
182 target_kind: ClosureKind,
183 ) -> Result<TraitImplRef, Error> {
184 self.translate_item(
185 span,
186 &closure.item,
187 TransItemSourceKind::TraitImpl(TraitImplSource::Closure(target_kind)),
188 )
189 }
190
191 pub fn translate_closure_info(
192 &mut self,
193 span: Span,
194 args: &hax::ClosureArgs,
195 ) -> Result<ClosureInfo, Error> {
196 use ClosureKind::*;
197 let kind = translate_closure_kind(&args.kind);
198
199 let fn_once_impl = self.translate_closure_bound_impl_ref(span, args, FnOnce)?;
200 let fn_mut_impl = if matches!(kind, FnMut | Fn) {
201 Some(self.translate_closure_bound_impl_ref(span, args, FnMut)?)
202 } else {
203 None
204 };
205 let fn_impl = if matches!(kind, Fn) {
206 Some(self.translate_closure_bound_impl_ref(span, args, Fn)?)
207 } else {
208 None
209 };
210 let signature = self.translate_poly_fun_sig(span, &args.fn_sig)?;
211 Ok(ClosureInfo {
212 kind,
213 fn_once_impl,
214 fn_mut_impl,
215 fn_impl,
216 signature,
217 })
218 }
219
220 pub fn get_closure_state_ty(
221 &mut self,
222 span: Span,
223 args: &hax::ClosureArgs,
224 ) -> Result<Ty, Error> {
225 let tref = self.translate_closure_type_ref(span, args)?;
226 Ok(TyKind::Adt(tref).into_ty())
227 }
228
229 pub fn translate_closure_upvar_tys(
233 &mut self,
234 span: Span,
235 args: &hax::ClosureArgs,
236 ) -> Result<IndexVec<FieldId, Ty>, Error> {
237 args.upvar_tys
238 .iter()
239 .map(|ty| self.translate_ty(span, ty))
240 .try_collect()
241 }
242
243 pub fn translate_closure_adt(
244 &mut self,
245 span: Span,
246 _args: &hax::ClosureArgs,
247 ) -> Result<TypeDeclKind, Error> {
248 let fields: IndexVec<FieldId, Field> = self
249 .the_only_binder()
250 .closure_upvar_tys
251 .as_ref()
252 .unwrap()
253 .iter()
254 .cloned()
255 .map(|ty| Field {
256 span,
257 attr_info: AttrInfo::dummy_private(),
258 name: None,
259 ty,
260 })
261 .collect();
262 Ok(TypeDeclKind::Struct(fields))
263 }
264
265 fn translate_closure_method_sig(
268 &mut self,
269 def: &hax::FullDef<'tcx>,
270 span: Span,
271 args: &hax::ClosureArgs,
272 target_kind: ClosureKind,
273 ) -> Result<RegionBinder<FunSig>, Error> {
274 let signature = &args.fn_sig;
275 trace!(
276 "signature of closure {:?}:\n{:?}",
277 def.def_id(),
278 signature.value,
279 );
280
281 let mut bound_regions = IndexVec::new();
282 let mut fun_sig = self
283 .translate_fun_sig(span, signature.hax_skip_binder_ref())?
284 .move_under_binder();
285 let state_ty = self.get_closure_state_ty(span, args)?.move_under_binder();
286
287 let state_ty = match target_kind {
289 ClosureKind::FnOnce => state_ty,
290 ClosureKind::Fn | ClosureKind::FnMut => {
291 let rid = bound_regions.push_with(|index| RegionParam::new(index, None));
292 let r = Region::Var(DeBruijnVar::new_at_zero(rid));
293 let mutability = if target_kind == ClosureKind::Fn {
294 RefKind::Shared
295 } else {
296 RefKind::Mut
297 };
298 TyKind::Ref(r, state_ty, mutability).into_ty()
299 }
300 };
301
302 let input_tys: Vec<Ty> = mem::take(&mut fun_sig.inputs);
304 fun_sig.inputs = vec![state_ty, Ty::mk_tuple(input_tys)];
306
307 Ok(RegionBinder {
308 regions: bound_regions,
309 skip_binder: fun_sig,
310 })
311 }
312
313 fn translate_closure_method_body(
314 &mut self,
315 span: Span,
316 def: &hax::FullDef<'tcx>,
317 target_kind: ClosureKind,
318 args: &hax::ClosureArgs,
319 signature: &FunSig,
320 ) -> Result<Body, Error> {
321 use ClosureKind::*;
322 let closure_kind = translate_closure_kind(&args.kind);
323 Ok(match (target_kind, closure_kind) {
324 (Fn, Fn) | (FnMut, FnMut) | (FnOnce, FnOnce) => {
325 let mut body = self.translate_def_body(span, def);
327 let Body::Unstructured(GExprBody {
336 locals,
337 body: blocks,
338 ..
339 }) = &mut body
340 else {
341 return Ok(body);
342 };
343
344 let tupled_ty = &signature.inputs[1];
346
347 blocks.dyn_visit_mut(|local: &mut LocalId| {
348 if local.index() >= 2 {
349 *local += 1;
350 }
351 });
352
353 let mut old_locals = mem::take(&mut locals.locals).into_iter();
354 locals.arg_count = 2;
355 locals.locals.push(old_locals.next().unwrap()); locals.locals.push(old_locals.next().unwrap()); let tupled_arg = locals.new_var(Some("tupled_args".to_string()), tupled_ty.clone());
358 locals.locals.extend(old_locals.map(|mut l| {
359 l.index += 1;
360 l
361 }));
362
363 let untupled_args = tupled_ty.as_tuple().unwrap();
364 let closure_arg_count = untupled_args.len();
365 let new_stts = untupled_args.iter().cloned().enumerate().map(|(i, ty)| {
366 let nth_field = tupled_arg.clone().project(
367 ProjectionElem::Field(
368 FieldProjKind::Tuple(closure_arg_count),
369 FieldId::new(i),
370 ),
371 ty,
372 );
373 let local_id = LocalId::new(i + 3);
374 Statement::new(
375 span,
376 StatementKind::Assign(
377 locals.place_for_var(local_id),
378 Rvalue::Use(Operand::Move(nth_field)),
379 ),
380 )
381 });
382 blocks[BlockId::ZERO].statements.splice(0..0, new_stts);
383
384 body
385 }
386 (FnOnce, Fn | FnMut) => {
396 let Some(body) = def.this.closure_once_shim(self.hax_state()) else {
398 panic!("missing shim for closure")
399 };
400 self.translate_body(span, body, &def.source_text)
401 }
402 (FnMut, Fn) => {
409 let fun_id: FunDeclId = self.register_item(
410 span,
411 def.this(),
412 TransItemSourceKind::ClosureMethod(closure_kind),
413 );
414 let impl_ref = self.translate_closure_impl_ref(span, args, closure_kind)?;
415 let fn_op = FnOperand::Regular(FnPtr::new(
418 fun_id.into(),
419 impl_ref.generics.concat(&GenericArgs {
420 regions: vec![self.translate_erased_region()].into(),
421 ..GenericArgs::empty()
422 }),
423 ));
424
425 let mut builder = BodyBuilder::new(span, 2);
426
427 let output = builder.new_var(None, signature.output.clone());
428 let state = builder.new_var(Some("state".to_string()), signature.inputs[0].clone());
429 let args = builder.new_var(Some("args".to_string()), signature.inputs[1].clone());
430 let deref_state = state.deref();
431 let reborrow_ty = TyKind::Ref(
432 self.translate_erased_region(),
433 deref_state.ty.clone(),
434 RefKind::Shared,
435 )
436 .into_ty();
437 let reborrow = builder.new_var(None, reborrow_ty);
438
439 builder.push_statement(StatementKind::Assign(
440 reborrow.clone(),
441 Rvalue::Ref {
442 place: deref_state,
443 kind: BorrowKind::Shared,
444 ptr_metadata: Operand::mk_const_unit(),
446 },
447 ));
448
449 builder.call(Call {
450 func: fn_op,
451 args: vec![Operand::Move(reborrow), Operand::Move(args)],
452 dest: output,
453 });
454
455 Body::Unstructured(builder.build())
456 }
457 (Fn, FnOnce) | (Fn, FnMut) | (FnMut, FnOnce) => {
458 panic!(
459 "Can't make a closure body for a more restrictive kind \
460 than the closure kind"
461 )
462 }
463 })
464 }
465
466 #[tracing::instrument(skip(self, item_meta))]
469 pub fn translate_closure_method(
470 mut self,
471 def_id: FunDeclId,
472 item_meta: ItemMeta,
473 def: &hax::FullDef<'tcx>,
474 target_kind: ClosureKind,
475 ) -> Result<FunDecl, Error> {
476 let span = item_meta.span;
477 let hax::FullDefKind::Closure {
478 args,
479 fn_once_impl,
480 fn_mut_impl,
481 fn_impl,
482 ..
483 } = &def.kind
484 else {
485 unreachable!()
486 };
487
488 let vimpl = match target_kind {
490 ClosureKind::FnOnce => fn_once_impl,
491 ClosureKind::FnMut => fn_mut_impl.as_ref().unwrap(),
492 ClosureKind::Fn => fn_impl.as_ref().unwrap(),
493 };
494 let implemented_trait = self.translate_trait_predicate(span, &vimpl.trait_pred)?;
495 let method_id = self.translate_trait_method_id(implemented_trait.id, &vimpl.methods[0])?;
496
497 let impl_ref = self.translate_closure_impl_ref(span, args, target_kind)?;
498 let src = ItemSource::TraitImpl {
499 impl_ref,
500 trait_ref: implemented_trait,
501 item_id: method_id.into(),
502 reuses_default: false,
503 };
504
505 let bound_sig = self.translate_closure_method_sig(def, span, args, target_kind)?;
507 let signature = bound_sig.apply(
509 self.the_only_binder()
510 .closure_call_method_region
511 .iter()
512 .map(|r| Region::Var(DeBruijnVar::new_at_zero(*r)))
513 .collect(),
514 );
515
516 let body = if item_meta.opacity.with_private_contents().is_opaque() {
517 Body::Opaque
518 } else {
519 self.translate_closure_method_body(span, def, target_kind, args, &signature)?
520 };
521
522 Ok(FunDecl {
523 def_id,
524 item_meta,
525 generics: self.into_generics(),
526 signature,
527 src,
528 is_global_initializer: None,
529 body,
530 })
531 }
532
533 #[tracing::instrument(skip(self, item_meta))]
534 pub fn translate_closure_trait_impl(
535 mut self,
536 def_id: TraitImplId,
537 item_meta: ItemMeta,
538 def: &hax::FullDef<'tcx>,
539 target_kind: ClosureKind,
540 ) -> Result<TraitImpl, Error> {
541 let span = item_meta.span;
542 let hax::FullDefKind::Closure {
543 args,
544 fn_once_impl,
545 fn_mut_impl,
546 fn_impl,
547 ..
548 } = def.kind()
549 else {
550 unreachable!()
551 };
552
553 let vimpl = match target_kind {
555 ClosureKind::FnOnce => fn_once_impl,
556 ClosureKind::FnMut => fn_mut_impl.as_ref().unwrap(),
557 ClosureKind::Fn => fn_impl.as_ref().unwrap(),
558 };
559 let mut timpl = self.translate_virtual_trait_impl(def_id, item_meta, vimpl)?;
560
561 if self.monomorphize() {
562 return Ok(timpl);
563 }
564
565 let trait_decl_id = timpl.impl_trait.id;
567 let trait_method_id = self.translate_trait_method_id(trait_decl_id, &vimpl.methods[0])?;
568 let call_fn_binder = {
569 let kind = TransItemSourceKind::ClosureMethod(target_kind);
570 let bound_method_ref: RegionBinder<DeclRef<ItemId>> =
571 self.translate_closure_bound_ref_with_method_bound(span, args, kind, target_kind)?;
572 let params = GenericParams {
573 regions: bound_method_ref.regions,
574 ..GenericParams::empty()
575 };
576 let fn_decl_ref: FunDeclRef = bound_method_ref.skip_binder.try_into().unwrap();
577 Binder::new(
578 BinderKind::TraitMethod(trait_decl_id, trait_method_id),
579 params,
580 fn_decl_ref,
581 )
582 };
583 timpl
584 .methods
585 .set_slot_extend(trait_method_id, call_fn_binder);
586
587 Ok(timpl)
588 }
589
590 #[tracing::instrument(skip(self, item_meta))]
593 pub fn translate_stateless_closure_as_fn(
594 mut self,
595 def_id: FunDeclId,
596 item_meta: ItemMeta,
597 def: &hax::FullDef<'tcx>,
598 ) -> Result<FunDecl, Error> {
599 let span = item_meta.span;
600 let hax::FullDefKind::Closure { args: closure, .. } = &def.kind else {
601 unreachable!()
602 };
603
604 trace!("About to translate closure as fn:\n{:?}", def.def_id());
605
606 assert!(
607 closure.upvar_tys.is_empty(),
608 "Only stateless closures can be translated as functions"
609 );
610
611 let signature = self.translate_fun_sig(span, closure.fn_sig.hax_skip_binder_ref())?;
613 let state_ty = self.get_closure_state_ty(span, closure)?;
614
615 let body = if item_meta.opacity.with_private_contents().is_opaque() {
616 Body::Opaque
617 } else {
618 let fun_id: FunDeclId = self.register_item(
626 span,
627 def.this(),
628 TransItemSourceKind::ClosureMethod(ClosureKind::FnOnce),
629 );
630 let impl_ref = self.translate_closure_impl_ref(span, closure, ClosureKind::FnOnce)?;
631 let fn_op = FnOperand::Regular(FnPtr::new(fun_id.into(), impl_ref.generics.clone()));
632
633 let mut builder = BodyBuilder::new(span, signature.inputs.len());
634
635 let output = builder.new_var(None, signature.output.clone());
636 let args: Vec<Place> = signature
637 .inputs
638 .iter()
639 .enumerate()
640 .map(|(i, ty)| builder.new_var(Some(format!("arg{}", i + 1)), ty.clone()))
641 .collect();
642 let args_tupled_ty = Ty::mk_tuple(signature.inputs.clone());
643 let args_tupled = builder.new_var(Some("args".to_string()), args_tupled_ty.clone());
644 let state = builder.new_var(Some("state".to_string()), state_ty.clone());
645
646 builder.push_statement(StatementKind::Assign(
647 args_tupled.clone(),
648 Rvalue::Aggregate(
649 AggregateKind::Adt(args_tupled_ty.as_adt().unwrap().clone(), None, None),
650 args.into_iter().map(Operand::Move).collect(),
651 ),
652 ));
653
654 let state_ty_adt = state_ty.as_adt().unwrap();
655 builder.push_statement(StatementKind::Assign(
656 state.clone(),
657 Rvalue::Aggregate(AggregateKind::Adt(state_ty_adt.clone(), None, None), vec![]),
658 ));
659
660 builder.call(Call {
661 func: fn_op,
662 args: vec![Operand::Move(state), Operand::Move(args_tupled)],
663 dest: output,
664 });
665
666 Body::Unstructured(builder.build())
667 };
668
669 Ok(FunDecl {
670 def_id,
671 item_meta,
672 generics: self.into_generics(),
673 signature,
674 src: ItemSource::TopLevel,
675 is_global_initializer: None,
676 body,
677 })
678 }
679}