add method to retrieve body of coroutine · rust-lang/rust@e40515a (original) (raw)

``

1

`+

//@ run-pass

`

``

2

`+

//! Tests stable mir API for retrieving the body of a coroutine.

`

``

3

+

``

4

`+

//@ ignore-stage1

`

``

5

`+

//@ ignore-cross-compile

`

``

6

`+

//@ ignore-remote

`

``

7

`+

//@ edition: 2024

`

``

8

+

``

9

`+

#![feature(rustc_private)]

`

``

10

`+

#![feature(assert_matches)]

`

``

11

+

``

12

`+

extern crate rustc_middle;

`

``

13

`+

#[macro_use]

`

``

14

`+

extern crate rustc_smir;

`

``

15

`+

extern crate rustc_driver;

`

``

16

`+

extern crate rustc_interface;

`

``

17

`+

extern crate stable_mir;

`

``

18

+

``

19

`+

use std::io::Write;

`

``

20

`+

use std::ops::ControlFlow;

`

``

21

+

``

22

`+

use stable_mir::mir::Body;

`

``

23

`+

use stable_mir::ty::{RigidTy, TyKind};

`

``

24

+

``

25

`+

const CRATE_NAME: &str = "crate_coroutine_body";

`

``

26

+

``

27

`+

fn test_coroutine_body() -> ControlFlow<()> {

`

``

28

`+

let crate_items = stable_mir::all_local_items();

`

``

29

`+

if let Some(body) = crate_items.iter().find_map(|item| {

`

``

30

`+

let item_ty = item.ty();

`

``

31

`+

if let TyKind::RigidTy(RigidTy::Coroutine(def, ..)) = &item_ty.kind() {

`

``

32

`+

if def.0.name() == "gbc::{closure#0}".to_string() {

`

``

33

`+

def.body()

`

``

34

`+

} else {

`

``

35

`+

None

`

``

36

`+

}

`

``

37

`+

} else {

`

``

38

`+

None

`

``

39

`+

}

`

``

40

`+

}) {

`

``

41

`+

check_coroutine_body(body);

`

``

42

`+

} else {

`

``

43

`` +

panic!("Cannot find gbc::{{closure#0}}. All local items are: {:#?}", crate_items);

``

``

44

`+

}

`

``

45

+

``

46

`+

ControlFlow::Continue(())

`

``

47

`+

}

`

``

48

+

``

49

`+

fn check_coroutine_body(body: Body) {

`

``

50

`+

let ret_ty = &body.locals()[0].ty;

`

``

51

`+

let local_3 = &body.locals()[3].ty;

`

``

52

`+

let local_4 = &body.locals()[4].ty;

`

``

53

+

``

54

`+

let TyKind::RigidTy(RigidTy::Adt(def, ..)) = &ret_ty.kind()

`

``

55

`+

else {

`

``

56

`+

panic!("Expected RigidTy::Adt, got: {:#?}", ret_ty);

`

``

57

`+

};

`

``

58

+

``

59

`+

assert_eq!("std::task::Poll", def.0.name());

`

``

60

+

``

61

`+

let TyKind::RigidTy(RigidTy::Coroutine(def, ..)) = &local_3.kind()

`

``

62

`+

else {

`

``

63

`+

panic!("Expected RigidTy::Coroutine, got: {:#?}", local_3);

`

``

64

`+

};

`

``

65

+

``

66

`+

assert_eq!("gbc::{closure#0}::{closure#0}", def.0.name());

`

``

67

+

``

68

`+

let TyKind::RigidTy(RigidTy::Coroutine(def, ..)) = &local_4.kind()

`

``

69

`+

else {

`

``

70

`+

panic!("Expected RigidTy::Coroutine, got: {:#?}", local_4);

`

``

71

`+

};

`

``

72

+

``

73

`+

assert_eq!("gbc::{closure#0}::{closure#0}", def.0.name());

`

``

74

`+

}

`

``

75

+

``

76

`+

fn main() {

`

``

77

`+

let path = "coroutine_body.rs";

`

``

78

`+

generate_input(&path).unwrap();

`

``

79

`+

let args = &[

`

``

80

`+

"rustc".to_string(),

`

``

81

`+

"-Cpanic=abort".to_string(),

`

``

82

`+

"--edition".to_string(),

`

``

83

`+

"2024".to_string(),

`

``

84

`+

"--crate-name".to_string(),

`

``

85

`+

CRATE_NAME.to_string(),

`

``

86

`+

path.to_string(),

`

``

87

`+

];

`

``

88

`+

run!(args, test_coroutine_body).unwrap();

`

``

89

`+

}

`

``

90

+

``

91

`+

fn generate_input(path: &str) -> std::io::Result<()> {

`

``

92

`+

let mut file = std::fs::File::create(path)?;

`

``

93

`+

write!(

`

``

94

`+

file,

`

``

95

`+

r#"

`

``

96

`+

async fn gbc() -> i32 {{

`

``

97

`+

let a = async {{ 1 }}.await;

`

``

98

`+

a

`

``

99

`+

}}

`

``

100

+

``

101

`+

fn main() {{}}

`

``

102

`+

"#

`

``

103

`+

)?;

`

``

104

`+

Ok(())

`

``

105

`+

}

`