1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

//! Provides C-like macros: offset_of and container_of
//!
//! 和memoffset crate的区别
//! Differences between crate memoffset and crate self
//! 1. 基于指针操作，不占用栈空间，支持超大数据结构，无栈溢出风险
//! 1. Pointer-based opertions do not occupy stack space, support super-large data structures, and avoid stack overflow risks.
//! 2. 提供container_of操作
//! 2. Provides the container_of operation.
//!
//! 基于指针的操作属于unsafe范围，如果使用声明宏方式，可能出现嵌套unsafe的编译告警，而过程宏可以消除这类告警。
//! Pointer-based operations are unsafe, If the declaration macro mode is used, nested unsafe compilation warnings may be generated, which can be eliminated by procedure macro.
//!
//! # Examples
//!
//! ```rust
//!
//!	extern crate hun_offsetof as hun;
//!
//! #[derive(Debug, Copy, Clone)]
//! #[repr(C)]
//! struct Bar {
//! 	key:	i32,
//!		value:	i32,
//! }
//!
//! #[derive(Debug)]
//! #[repr(C)]
//! struct Foo {
//!		key:	i32,
//!		value:	[Bar; 2],
//! }
//!
//!	assert_eq!(hun::offset_of!(Bar, value), 4);
//! assert_eq!(hun::offset_of!(Foo, value[1].key), 12);
//!
//! let foo = Foo {
//!		key: 1,
//!		value: [ Bar { key: 2, value: 2}, Bar { key: 3, value: 3 }],
//!	};
//!	let value = &foo.value[1].value;
//!
//! let obj = unsafe { hun::container_of!(value, Foo, value[1].value) };
//!	assert_eq!(obj as *const _, &foo as *const _);
//! ```
//!

extern crate proc_macro;
use std::time::{ SystemTime };
use proc_macro::{ TokenStream };
use proc_macro2::{ Span };
use quote::quote;
use syn::{
	Expr, Token, Path, parse_macro_input,
	spanned::{ Spanned },
	parse::{self, Parse, ParseStream },
	punctuated::{ Punctuated },
};

#[proc_macro]
pub fn offset_of(input: TokenStream) -> TokenStream {
	let input = parse_macro_input!(input as OffsetOfInput);
	let parent = &input.parent;
	let member = &input.member;
	let base = make_ident("base", parent.span());
	quote! {
		unsafe {
			let #base = 0 as *const #parent;
			&((*#base).#member) as *const _ as *const u8 as usize
		}
	}.into()
}

#[proc_macro]
pub fn container_of(input: TokenStream) -> TokenStream {
	container_of_impl(input, false)
}

#[proc_macro]
pub fn container_of_mut(input: TokenStream) -> TokenStream {
	container_of_impl(input, true)
}

struct OffsetOfInput {
	parent: Path,
	_comma: Option<Token![,]>,
	member:	Punctuated<Expr, Token![.]>,
}

impl Parse for OffsetOfInput {
	fn parse(input: ParseStream) -> parse::Result<Self> {
		Ok(OffsetOfInput {
			parent: input.parse()?,
			_comma: input.parse()?,
			member: Punctuated::parse_terminated(input)?,
		})
	}
}

fn container_of_impl(input: TokenStream, is_mut: bool) -> TokenStream {
	let input = parse_macro_input!(input as ContainerOfInput);
	let tokens = container_of_tokens(&input);
	if is_mut {
		quote!({ &mut *(#tokens) }).into()
	} else {
		quote!({ &*(#tokens) }).into()
	}
}

fn container_of_tokens(input: &ContainerOfInput) -> proc_macro2::TokenStream {
	let obj = &input.obj;
	let parent = &input.offset.parent;
	let member = &input.offset.member;
	let base = make_ident("base", parent.span());
	let offset = make_ident("offset", member.span());
	quote!({
		let #base = 0 as *const #parent;
		let #offset = &((*#base).#member) as *const _ as *const u8 as usize;
		if ::core::mem::size_of_val(&((*#base).#member)) ==
			::core::mem::size_of_val(#obj) {
			let #base = (#obj) as *const _ as *const u8 as usize;
			(#base - #offset) as *mut u8 as *mut #parent
		} else {
			core::panic!("container_of: wrong arguments!");
		}
	})
}

struct ContainerOfInput {
	obj:	Expr,
	_comma:	Option<Token![,]>,
	offset:	OffsetOfInput,
}

impl Parse for ContainerOfInput {
	fn parse(input: ParseStream) -> syn::Result<Self> {
		Ok(ContainerOfInput {
			obj:	input.parse()?,
			_comma:	input.parse()?,
			offset:	OffsetOfInput {
				parent:	input.parse()?,
				_comma:	input.parse()?,
				member:	Punctuated::parse_terminated(input)?,
			},
		})
	}
}

fn make_ident(name: &str, span: Span) -> syn::Ident {
	let name = format!("{}_{}", name, SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs());
	syn::Ident::new(&name, span)
}