Skip to content

Use a type visitor to collect type metadata, remove layout #68

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 16 commits into from
Apr 29, 2025
Merged

Use a type visitor to collect type metadata, remove layout #68

Show file tree
Hide file tree
Changes from 13 commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
f0817f1
TyVisitor which uses stable_mir::Visitor
gtrepta Apr 25, 2025
fe07f42
Use collected types from TyVisitor
gtrepta Apr 18, 2025
4e6d962
Remove dead code
gtrepta Apr 25, 2025
d5b28a6
Rename TyVisitor -> TyCollector
gtrepta Apr 25, 2025
2e27d89
Update golden tests.
gtrepta Apr 25, 2025
5d8fc46
Use "None" for layout
gtrepta Apr 22, 2025
2eae9c6
Revert "Use "None" for layout"
gtrepta Apr 22, 2025
af75f9f
Resolve FnDefs
gtrepta Apr 22, 2025
905a5c2
Resolve closures. Check for types that have already been visited.
gtrepta Apr 25, 2025
e161f38
Check for already resolved tys (prevent infinite loops)
gtrepta Apr 25, 2025
cdb1be6
Visit bottom up. Break on what we don't support yet.
gtrepta Apr 25, 2025
304e0eb
Break on CoroutineWitnesses, because we can't grab the layout on them.
gtrepta Apr 25, 2025
f403c59
update golden tests
gtrepta Apr 25, 2025
6919e37
Handling `FnPtr` signature internally to avoid `Binder`s (#71)
dkcumming Apr 26, 2025
e652b72
make golden
dkcumming Apr 27, 2025
9262074
Improve normalizing filter
gtrepta Apr 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
197 changes: 90 additions & 107 deletions src/printer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,13 @@
use std::any::Any;
use std::io::Write;
use std::{collections::HashMap, fs::File, io, iter::Iterator, str, vec::Vec};
use std::ops::ControlFlow;
use std::{
collections::{HashMap, HashSet},
fs::File,
io,
iter::Iterator,
str,
vec::Vec,
};
extern crate rustc_middle;
extern crate rustc_monomorphize;
extern crate rustc_session;
Expand All @@ -25,9 +32,10 @@ use stable_mir::{
mir::mono::{Instance, InstanceKind, MonoItem},
mir::{alloc::AllocId, visit::MirVisitor, Body, LocalDecl, Rvalue, Terminator, TerminatorKind},
ty::{
AdtDef, Allocation, ConstDef, ForeignItemKind, IndexedVal, RigidTy, TyConstKind, TyKind,
VariantIdx,
AdtDef, Allocation, ConstDef, ForeignItemKind, IndexedVal, Region, RegionKind, RigidTy,
TyKind, VariantIdx,
},
visitor::{Visitable, Visitor},
CrateDef, CrateItem, ItemKind,
};

Expand Down Expand Up @@ -480,13 +488,84 @@ type AllocMap = HashMap<stable_mir::mir::alloc::AllocId, AllocInfo>;
type TyMap =
HashMap<stable_mir::ty::Ty, (stable_mir::ty::TyKind, Option<stable_mir::abi::LayoutShape>)>;

struct TyCollector {
types: TyMap,
resolved: HashSet<stable_mir::ty::Ty>,
}

impl TyCollector {
fn new() -> TyCollector {
TyCollector {
types: HashMap::new(),
resolved: HashSet::new(),
}
}
}

impl TyCollector {
#[inline(always)]
fn visit_instance(&mut self, instance: Instance) -> ControlFlow<<Self as Visitor>::Break> {
let fn_abi = instance.fn_abi().unwrap();
let mut inputs_outputs: Vec<stable_mir::ty::Ty> =
fn_abi.args.iter().map(|arg_abi| arg_abi.ty).collect();
inputs_outputs.push(fn_abi.ret.ty);
inputs_outputs.super_visit(self)
}
}

impl Visitor for TyCollector {
type Break = ();

fn visit_ty(&mut self, ty: &stable_mir::ty::Ty) -> ControlFlow<Self::Break> {
if self.types.contains_key(ty) || self.resolved.contains(ty) {
return ControlFlow::Continue(());
}

match ty.kind() {
TyKind::RigidTy(RigidTy::Closure(def, ref args)) => {
self.resolved.insert(*ty);
let instance =
Instance::resolve_closure(def, args, stable_mir::ty::ClosureKind::Fn).unwrap();
self.visit_instance(instance)
}
// Break on CoroutineWitnesses, because they aren't expected when getting the layout
TyKind::RigidTy(RigidTy::CoroutineWitness(..)) => ControlFlow::Break(()),
TyKind::RigidTy(RigidTy::FnDef(def, ref args)) => {
self.resolved.insert(*ty);
let instance = Instance::resolve(def, args).unwrap();
self.visit_instance(instance)
}
_ => {
let control = ty.super_visit(self);
match control {
ControlFlow::Continue(_) => {
let maybe_layout_shape = ty.layout().ok().map(|layout| layout.shape());
self.types.insert(*ty, (ty.kind(), maybe_layout_shape));
control
}
_ => control,
}
}
}
}

fn visit_reg(&mut self, reg: &Region) -> ControlFlow<Self::Break> {
match reg.kind {
// Breaking on Bound regions, because some have escaping bound vars
// which can't be wrapped in dummy binders.
RegionKind::ReBound(..) => ControlFlow::Break(()),
_ => ControlFlow::Continue(()),
}
}
}

struct InternedValueCollector<'tcx, 'local> {
tcx: TyCtxt<'tcx>,
_sym: String,
locals: &'local [LocalDecl],
link_map: &'local mut LinkMap<'tcx>,
visited_allocs: &'local mut AllocMap,
visited_tys: &'local mut TyMap,
ty_visitor: &'local mut TyCollector,
}

type InternedValues<'tcx> = (LinkMap<'tcx>, AllocMap, TyMap);
Expand Down Expand Up @@ -601,103 +680,6 @@ fn collect_alloc(
};
}

fn collect_vec_tys(collector: &mut InternedValueCollector, tys: Vec<stable_mir::ty::Ty>) {
tys.into_iter().for_each(|ty| collect_ty(collector, ty));
}

fn collect_arg_tys(collector: &mut InternedValueCollector, args: &stable_mir::ty::GenericArgs) {
use stable_mir::ty::{GenericArgKind::*, TyConstKind::*}; // TyConst
for arg in args.0.iter() {
match arg {
Type(ty) => collect_ty(collector, *ty),
Const(ty_const) => match ty_const.kind() {
Value(ty, _) | ZSTValue(ty) => collect_ty(collector, *ty),
_ => {}
},
_ => {}
}
}
}

fn collect_ty(val_collector: &mut InternedValueCollector, val: stable_mir::ty::Ty) {
use stable_mir::ty::{RigidTy::*, TyKind::RigidTy};

let maybe_layout_shape = val.layout().ok().map(|layout| layout.shape());

if val_collector
.visited_tys
.insert(val, (val.kind(), maybe_layout_shape))
.is_some()
{
match val.kind() {
RigidTy(Array(ty, ty_const)) => {
collect_ty(val_collector, ty);
match ty_const.kind() {
TyConstKind::Value(ty, _) => collect_ty(val_collector, *ty),
TyConstKind::ZSTValue(ty) => collect_ty(val_collector, *ty),
_ => (),
}
}
RigidTy(Pat(ty, _) | Slice(ty) | RawPtr(ty, _) | Ref(_, ty, _)) => {
collect_ty(val_collector, ty)
}
RigidTy(Tuple(tys)) => collect_vec_tys(val_collector, tys),
RigidTy(Adt(def, ref args)) => {
for variant in def.variants_iter() {
for field in variant.fields() {
collect_ty(val_collector, field.ty());
}
}
collect_arg_tys(val_collector, args);
}
// FIXME: Would be good to grab the coroutine signature
RigidTy(Coroutine(_, ref args, _) | CoroutineWitness(_, ref args)) => {
collect_arg_tys(val_collector, args)
}
RigidTy(Closure(def, ref args)) => {
// Need to monomorphise again
// let ty_internal = rustc_internal::internal(val_collector.tcx, val);
// let closure_kind = ty_internal.to_opt_closure_kind().unwrap(); // Errors on span for some reason
// Choose FnOnce, got the idea from kani reachability.rs
let instance =
Instance::resolve_closure(def, args, stable_mir::ty::ClosureKind::Fn).unwrap();
let fn_abi = instance.fn_abi().unwrap();
let mut inputs_and_outputs: Vec<stable_mir::ty::Ty> =
fn_abi.args.iter().map(|arg| arg.ty).collect();
if inputs_and_outputs[0].type_id() == val.type_id() {
// TODO: Not sure if it always occurs, but had case where the first ArgAbi had ty same as val.
// This causes an infinite loop
inputs_and_outputs.remove(0);
}
inputs_and_outputs.push(fn_abi.ret.ty);

collect_vec_tys(val_collector, inputs_and_outputs);
collect_arg_tys(val_collector, args); // Do I need to do this again?
}
ref _kind @ RigidTy(FnDef(def, ref args)) => {
// Need to monomorphise again
let instance = Instance::resolve(def, args).unwrap();
let fn_abi = instance.fn_abi().unwrap();
let mut inputs_and_outputs: Vec<stable_mir::ty::Ty> =
fn_abi.args.iter().map(|arg| arg.ty).collect();
inputs_and_outputs.push(fn_abi.ret.ty);

collect_vec_tys(val_collector, inputs_and_outputs);
collect_arg_tys(val_collector, args); // Do I need to do this again?
}
RigidTy(Foreign(def)) => match def.kind() {
ForeignItemKind::Fn(def) => {
// I think this is fully monomorphised???
collect_vec_tys(val_collector, def.fn_sig().value.inputs_and_output)
}
ForeignItemKind::Type(ty) => collect_ty(val_collector, ty),
ForeignItemKind::Static(def) => collect_ty(val_collector, def.ty()),
},
_ => {} // TODO: I think we need to add arm for FnPtr? Maybe we should create an example that uses it first
}
}
}

impl MirVisitor for InternedValueCollector<'_, '_> {
fn visit_terminator(&mut self, term: &Terminator, loc: stable_mir::mir::visit::Location) {
use stable_mir::mir::{ConstOperand, Operand::Constant};
Expand Down Expand Up @@ -772,14 +754,15 @@ impl MirVisitor for InternedValueCollector<'_, '_> {
}

fn visit_ty(&mut self, ty: &stable_mir::ty::Ty, _location: stable_mir::mir::visit::Location) {
collect_ty(self, *ty);
ty.visit(self.ty_visitor);
self.super_ty(ty);
}
}

fn collect_interned_values<'tcx>(tcx: TyCtxt<'tcx>, items: Vec<&MonoItem>) -> InternedValues<'tcx> {
let mut calls_map = HashMap::new();
let mut visited_tys = HashMap::new();
let mut visited_allocs = HashMap::new();
let mut ty_visitor = TyCollector::new();
if link_items_enabled() {
for item in items.iter() {
if let MonoItem::Fn(inst) = item {
Expand All @@ -800,8 +783,8 @@ fn collect_interned_values<'tcx>(tcx: TyCtxt<'tcx>, items: Vec<&MonoItem>) -> In
_sym: inst.mangled_name(),
locals: body.locals(),
link_map: &mut calls_map,
visited_tys: &mut visited_tys,
visited_allocs: &mut visited_allocs,
ty_visitor: &mut ty_visitor,
}
.visit_body(&body)
} else {
Expand All @@ -819,8 +802,8 @@ fn collect_interned_values<'tcx>(tcx: TyCtxt<'tcx>, items: Vec<&MonoItem>) -> In
_sym: inst.mangled_name(),
locals: &[],
link_map: &mut calls_map,
visited_tys: &mut visited_tys,
visited_allocs: &mut visited_allocs,
ty_visitor: &mut ty_visitor,
}
.visit_body(&body)
} else {
Expand All @@ -833,7 +816,7 @@ fn collect_interned_values<'tcx>(tcx: TyCtxt<'tcx>, items: Vec<&MonoItem>) -> In
MonoItem::GlobalAsm(_) => {}
}
}
(calls_map, visited_allocs, visited_tys)
(calls_map, visited_allocs, ty_visitor.types)
}

// Collection Transitive Closure
Expand Down
86 changes: 7 additions & 79 deletions tests/integration/programs/assert_eq.smir.json.expected
View file
Open in desktop
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am a bit puzzled by the differences in the expectation files.
The code does not collect any function types any more ? That's OK I guess, we don't use them, but... we also don't get I8 any more (in this example)? Why not?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I also can't understand why i8 is now removed from the map

Copy link
Contributor

@gtrepta gtrepta Apr 28, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed functions aren't in the map as well. I think it's because they get inserted by Stephen's code before monomorphization:

    if val_collector
.visited_tys
.insert(val, (val.kind(), maybe_layout_shape))
.is_some()
{
match val.kind() {

The monomorphization gets done inside the match statement, but the function is already added to the map! In my code, the argument and return types for the function are added only. But, we should probably figure out how to get the resolved function signature and add it as well.

Not sure why the I8 disappears, idk why it's there in the first place either. The integers in the source don't have any type annotations, rust defaults to I32 (which is in the type map), so there shouldn't be any issue with that.

Original file line number Diff line number Diff line change
Expand Up @@ -3119,13 +3119,6 @@
}
],
"types": [
[
{
"PrimitiveType": {
"Int": "I8"
}
}
],
[
{
"PrimitiveType": {
Expand Down Expand Up @@ -3255,21 +3248,21 @@
[
{
"StructType": {
"name": "std::fmt::Error"
"name": "std::fmt::Formatter"
}
}
],
[
{
"StructType": {
"name": "std::fmt::Formatter"
"name": "std::fmt::Error"
}
}
],
[
{
"TupleType": {
"types": "elided"
"StructType": {
"name": "std::panic::Location"
}
}
],
Expand Down Expand Up @@ -3311,7 +3304,7 @@
],
[
{
"RefType": "elided"
"PtrType": "elided"
}
],
[
Expand Down Expand Up @@ -3351,77 +3344,12 @@
],
[
{
"FunType": "elided"
}
],
[
{
"FunType": "elided"
}
],
[
{
"FunType": "elided"
}
],
[
{
"FunType": "elided"
}
],
[
{
"FunType": "elided"
}
],
[
{
"FunType": "elided"
}
],
[
{
"FunType": "elided"
}
],
[
{
"FunType": "elided"
}
],
[
{
"FunType": "elided"
}
],
[
{
"FunType": "elided"
}
],
[
{
"FunType": "elided"
}
],
[
{
"FunType": "elided"
}
],
[
{
"FunType": "elided"
}
],
[
{
"FunType": "elided"
"RefType": "elided"
}
],
[
{
"FunType": "elided"
"RefType": "elided"
}
],
[
Expand Down
Loading
Loading