portal/pit
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

stuff

Browse source
This commit is contained in:
Graham Kelly 2024-06-24 07:21:05 -04:00
commit 7e642e1f85
16 changed files with 2504 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored Normal file
View file

@ -0,0 +1 @@
target

1024
Cargo.lock generated Normal file
View file

File diff suppressed because it is too large Load diff

7
Cargo.toml Normal file
View file

@ -0,0 +1,7 @@
[workspace]
members=["pit-core", "pit-patch", "pit-rust-guest", "pit-rust-host", "pit-rust-host-lib"]
resolver="2"
[workspace.dependencies]
portal-pc-waffle = {git="https://github.com/portal-co/waffle-.git",branch="pr/changes2"}
waffle-ast = {git="https://github.com/portal-co/more_waffle.git",branch="master"}

12
pit-core/Cargo.toml Normal file
View file

@ -0,0 +1,12 @@
[package]
name = "pit-core"
version = "0.1.0"
edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
derive_more = "0.99.18"
hex = "0.4.3"
nom = "7.1.3"
sha3 = "0.10.8"

157
pit-core/src/lib.rs Normal file
View file

@ -0,0 +1,157 @@
use derive_more::Display;
use nom::{
bytes::complete::{is_not, tag, take, take_while_m_n},
character::complete::{alpha1, char, multispace0},
combinator::opt,
multi::{many0, separated_list0},
sequence::{delimited, tuple},
IResult,
};
use sha3::{Digest, Sha3_256};
use std::{collections::BTreeMap, fmt::Display};
#[derive(Display, Clone, Eq, PartialEq, PartialOrd, Ord)]
pub enum ResTy {
#[display(fmt = "")]
None,
#[display(fmt = "{}", "hex::encode(_0)")]
Of([u8; 32]),
#[display(fmt = "this")]
This,
}
pub fn parse_resty(a: &str) -> IResult<&str, ResTy> {
if let Some(a) = a.strip_prefix("this") {
// let (a, k) = opt(tag("n"))(a)?;
return Ok((a, ResTy::This));
}
let (a, d) = opt(take_while_m_n(64, 64, |a: char| a.is_digit(16)))(a)?;
return Ok((
a,
match d {
Some(d) => {
let mut b = [0u8; 32];
hex::decode_to_slice(d, &mut b).unwrap();
ResTy::Of(b)
}
None => ResTy::None,
},
));
}
#[derive(Display, Clone, Eq, PartialEq, PartialOrd, Ord)]
pub enum Arg {
I32,
I64,
F32,
F64,
#[display(fmt = "R{}{}{}", "ty", "if *nullable{\"n\"}else{\"\"}","if *take{\"\"}else{\"&\"}")]
Resource {
ty: ResTy,
nullable: bool,
take: bool,
},
// #[display(fmt = "{}", _0)]
// Func(Sig),
}
pub fn parse_arg(a: &str) -> IResult<&str, Arg> {
let (_, a) = multispace0(a)?;
let (b, c) = take(1usize)(a)?;
match c {
"R" => {
// if let Some(a) = b.strip_prefix("this"){
// let (a, k) = opt(tag("n"))(a)?;
// return Ok((
// a,
// Arg::Resource {
// ty: ResTy::This,
// nullable: k.is_some(),
// },
// ));
// }
let (a, d) = parse_resty(b)?;
let (a, k) = opt(tag("n"))(a)?;
let (a, take) = opt(tag("&"))(a)?;
return Ok((
a,
Arg::Resource {
ty: d,
nullable: k.is_some(),
take: take.is_none(),
},
));
}
// "(" => {
// let (a, x) = parse_sig(a)?;
// return Ok((a, Arg::Func(x)));
// }
_ => {
let (c, a) = take(3usize)(a)?;
match c {
"I32" => return Ok((a, Arg::I32)),
"I64" => return Ok((a, Arg::I64)),
"F32" => return Ok((a, Arg::F32)),
"F64" => return Ok((a, Arg::F64)),
_ => {}
}
}
}
todo!()
}
#[derive(Display, Clone, Eq, PartialEq, PartialOrd, Ord)]
#[display(
fmt = "({}) -> ({})",
"params.iter().map(|a|a.to_string()).collect::<Vec<_>>().join(\",\")",
"rets.iter().map(|a|a.to_string()).collect::<Vec<_>>().join(\",\")"
)]
pub struct Sig {
pub params: Vec<Arg>,
pub rets: Vec<Arg>,
}
pub fn parse_sig(a: &str) -> IResult<&str, Sig> {
let (a, _) = multispace0(a)?;
let mut d = delimited(char('('), separated_list0(char(','), parse_arg), char(')'));
let (a, params) = d(a)?;
let (a, _) = multispace0(a)?;
let (a, _) = tag("->")(a)?;
let (a, _) = multispace0(a)?;
let (a, rets) = d(a)?;
return Ok((a, Sig { params, rets }));
}
#[derive(Clone, Eq, PartialEq, PartialOrd, Ord)]
pub struct Interface {
pub methods: BTreeMap<String, Sig>,
}
impl Display for Interface {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{{")?;
for (i, (a, b)) in self.methods.iter().enumerate() {
if i != 0 {
write!(f, ";")?;
}
write!(f, "{}{}", a, b)?;
}
return write!(f, "}}");
}
}
pub fn parse_interface(a: &str) -> IResult<&str, Interface> {
pub fn go(a: &str) -> IResult<&str, Interface> {
let (a, s) = separated_list0(char(';'), tuple((multispace0, alpha1, parse_sig)))(a)?;
return Ok((
a,
Interface {
methods: s.into_iter().map(|(_, a, b)| (a.to_owned(), b)).collect(),
},
));
}
let (a, _) = multispace0(a)?;
return delimited(char('{'), go, char('}'))(a);
}
impl Interface {
pub fn rid(&self) -> [u8; 32] {
use std::io::Write;
let mut s = Sha3_256::default();
write!(s, "{}", self);
return s.finalize().try_into().unwrap();
}
pub fn rid_str(&self) -> String {
return hex::encode(self.rid());
}
}

13
pit-patch/Cargo.toml Normal file
View file

@ -0,0 +1,13 @@
[package]
name = "pit-patch"
version = "0.1.0"
edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
anyhow = "1.0.86"
nom = "7.1.3"
pit-core = { version = "0.1.0", path = "../pit-core" }
portal-pc-waffle.workspace = true
waffle-ast.workspace = true

164
pit-patch/src/canon.rs Normal file
View file

@ -0,0 +1,164 @@
use std::{collections::BTreeMap, iter::once, mem::take};
use anyhow::Context;
use waffle::{
entity::EntityRef, util::new_sig, BlockTarget, Export, ExportKind, Func, FuncDecl,
FunctionBody, Import, ImportKind, Module, Operator, SignatureData, Type,
};
use waffle_ast::{Builder, Expr};
pub fn canon(m: &mut Module, rid: &str, target: &str) -> anyhow::Result<()> {
let mut xs = vec![];
for i in m.imports.iter() {
if i.module == format!("pit/{rid}") {
if let Some(a) = i.name.strip_prefix("~") {
xs.push(a.to_owned())
}
}
}
let s = new_sig(
m,
SignatureData {
params: vec![Type::I32],
returns: vec![Type::ExternRef],
},
);
let f2 = m
.funcs
.push(waffle::FuncDecl::Import(s, format!("pit/{rid}.~{target}")));
for i in take(&mut m.imports) {
if i.module == format!("pit/{rid}") {
if let Some(a) = i.name.strip_prefix("~") {
if let Ok(x) = xs.binary_search(&a.to_owned()) {
if let ImportKind::Func(f) = i.kind {
let fs = m.funcs[f].sig();
let fname = m.funcs[f].name().to_owned();
let mut b = FunctionBody::new(&m, fs);
let k = b.entry;
let mut e = Expr::Bind(
Operator::I32Add,
vec![
Expr::Bind(Operator::I32Const { value: x as u32 }, vec![]),
Expr::Bind(
Operator::I32Mul,
vec![
Expr::Bind(
Operator::I32Const {
value: xs.len() as u32,
},
vec![],
),
Expr::Leaf(b.blocks[k].params[0].1),
],
),
],
);
let (a, k) = e.build(m, &mut b, k)?;
let args = once(a)
.chain(b.blocks[b.entry].params[1..].iter().map(|a| a.1))
.collect();
b.set_terminator(k, waffle::Terminator::ReturnCall { func: f2, args });
m.funcs[f] = FuncDecl::Body(fs, fname, b);
continue;
}
}
}
}
m.imports.push(i)
}
m.imports.push(Import {
module: format!("pit/{rid}"),
name: format!("~{target}"),
kind: ImportKind::Func(f2),
});
let mut b = BTreeMap::new();
for x in take(&mut m.exports) {
for x2 in xs.iter() {
if let Some(a) = x.name.strip_prefix(&format!("pit/{rid}/~{x2}")) {
let mut b = b.entry(a.to_owned()).or_insert_with(|| BTreeMap::new());
let mut e = b.entry(x2.clone()).or_insert_with(|| Func::invalid());
if let ExportKind::Func(f) = x.kind {
*e = f;
continue;
}
}
}
m.exports.push(x)
}
for (method, inner) in b.into_iter() {
let sig = *inner.iter().next().context("in getting an instance")?.1;
let funcs: Vec<Func> = xs
.iter()
.filter_map(|f| inner.get(f))
.cloned()
.collect::<Vec<_>>();
let sig = m.funcs[sig].sig();
let mut b = FunctionBody::new(&m, sig);
let k = b.entry;
let mut e = Expr::Bind(
Operator::I32DivU,
vec![
Expr::Leaf(b.blocks[k].params[0].1),
Expr::Bind(
Operator::I32Const {
value: funcs.len() as u32,
},
vec![],
),
],
);
let (a, k) = e.build(m, &mut b, k)?;
let mut e = Expr::Bind(
Operator::I32RemU,
vec![
Expr::Leaf(b.blocks[k].params[0].1),
Expr::Bind(
Operator::I32Const {
value: funcs.len() as u32,
},
vec![],
),
],
);
let (c, k) = e.build(m, &mut b, k)?;
let args = once(a)
.chain(b.blocks[b.entry].params[1..].iter().map(|a| a.1))
.collect::<Vec<_>>();
let blocks = funcs
.iter()
.map(|f| {
let k = b.add_block();
b.set_terminator(
k,
waffle::Terminator::ReturnCall {
func: *f,
args: args.clone(),
},
);
BlockTarget {
block: k,
args: vec![],
}
})
.collect::<Vec<_>>();
b.set_terminator(
k,
waffle::Terminator::Select {
value: c,
targets: blocks,
default: BlockTarget {
block: b.entry,
args,
},
},
);
let f = m
.funcs
.push(FuncDecl::Body(sig, format!("pit/{rid}/~{target}{method}"), b));
m.exports.push(Export {
name: format!("pit/{rid}/~{target}{method}"),
kind: ExportKind::Func(f),
});
}
Ok(())
}

33
pit-patch/src/lib.rs Normal file
View file

@ -0,0 +1,33 @@
use std::{
collections::{BTreeMap, BTreeSet},
iter::once,
mem::take,
};
use anyhow::Context;
use pit_core::{Arg, Interface, Sig};
use util::tfree;
use waffle::{
util::new_sig, ExportKind, FuncDecl, FunctionBody, ImportKind, Module, SignatureData,
TableData, Type,
};
use waffle_ast::add_op;
pub mod util;
pub mod canon;
pub mod lower;
pub fn get_interfaces(m: &Module) -> anyhow::Result<Vec<Interface>> {
let c = m
.custom_sections
.get(".pit-types")
.context("in getting type section")?;
let mut is = vec![];
for b in c.split(|a| *a == 0) {
let s = std::str::from_utf8(b)?;
let (s, i) = pit_core::parse_interface(s)
.map_err(|e: nom::Err<nom::error::Error<&str>>| anyhow::anyhow!("invalid pit"))?;
is.push(i);
}
return Ok(is);
}

309
pit-patch/src/lower.rs Normal file
View file

@ -0,0 +1,309 @@
use std::{collections::BTreeSet, iter::once, mem::take};
use anyhow::Context;
use waffle::{
entity::EntityRef, util::new_sig, ExportKind, Func, FuncDecl, FunctionBody, Import, ImportKind,
Module, Operator, SignatureData, TableData, Type, ValueDef,
};
use waffle_ast::{
add_op,
fcopy::{obf_mod, DontObf, Obfuscate},
Builder, Expr,
};
use crate::canon::canon;
pub fn patch_ty(t: &mut Type) {
if let Type::ExternRef = t.clone() {
*t = Type::I32
}
}
pub fn import_fn(m: &mut Module, mo: String, n: String, s: SignatureData) -> Func {
for i in m.imports.iter() {
if i.module == mo && i.name == n {
if let ImportKind::Func(f) = &i.kind {
return *f;
}
}
}
let s = new_sig(m, s);
let f = m
.funcs
.push(waffle::FuncDecl::Import(s, format!("{mo}.{n}")));
m.imports.push(Import {
module: mo,
name: n,
kind: ImportKind::Func(f),
});
return f;
}
pub fn instantiate(m: &mut Module, root: &str) -> anyhow::Result<()> {
let i = crate::get_interfaces(m)?;
let interfaces = i
.into_iter()
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
for j in interfaces.iter() {
canon(m, &j.rid_str(), root)?;
}
for s in m.signatures.values_mut() {
for p in s.params.iter_mut().chain(s.returns.iter_mut()) {
patch_ty(p)
}
}
m.try_take_per_func_body(|m, b| {
let x = b.type_pool.from_iter(once(Type::I32));
for p in b
.blocks
.values_mut()
.flat_map(|a| a.params.iter_mut().map(|a| &mut a.0))
.chain(b.rets.iter_mut())
{
patch_ty(p)
}
for v in b.values.iter().collect::<Vec<_>>() {
if let ValueDef::Operator(o, _1, tys) = &mut b.values[v] {
match o.clone() {
Operator::RefNull { ty } => {
if ty == Type::ExternRef {
*o = Operator::I32Const { value: 0 }
}
}
Operator::RefIsNull => {
if b.type_pool[*tys][0] == Type::ExternRef {
*o = Operator::I32Eqz
}
}
_ => {}
}
}
}
for t in b.type_pool.storage.iter_mut() {
patch_ty(t)
}
Ok::<_, anyhow::Error>(())
})?;
struct X {}
impl Obfuscate for X {
fn obf(
&mut self,
o: Operator,
f: &mut waffle::FunctionBody,
b: waffle::Block,
args: &[waffle::Value],
types: &[Type],
module: &mut Module,
) -> anyhow::Result<(waffle::Value, waffle::Block)> {
match o {
Operator::TableGet { table_index }
| Operator::TableSet { table_index }
| Operator::TableSize { table_index }
| Operator::TableGrow { table_index } => {
if module.tables[table_index].ty == Type::ExternRef {
// let x = b.arg_pool[*_1].to_vec();
let w = add_op(
f,
&[],
&[Type::I32],
Operator::I32Const {
value: table_index.index() as u32,
},
);
f.append_to_block(b, w);
let id = format!("pit-patch-rt/{}", o.to_string().split_once("<").unwrap().0);
let mut a = module.exports.iter();
let a = loop{
let Some(b) = a.next() else{
anyhow::bail!("pit patch rt not found")
};
if b.name != id{
continue;
}
let ExportKind::Func(a) = &b.kind else{
continue;
};
break *a;
};
DontObf {}.obf(
Operator::Call { function_index: a },
f,
b,
&once(w).chain(args.iter().cloned()).collect::<Vec<_>>(),
types,
module,
)
} else {
DontObf {}.obf(o, f, b, args, types, module)
}
}
_ => DontObf {}.obf(o, f, b, args, types, module),
}
}
}
obf_mod(m, &mut X {})?;
for t in m.tables.values_mut() {
if t.ty == Type::ExternRef {
t.ty = Type::FuncRef;
}
}
for i in take(&mut m.imports) {
if let Some(rid) = i.module.strip_prefix("pit/") {
let ridx = interfaces
.iter()
.enumerate()
.find_map(|x| {
if x.1.rid_str() == rid {
Some(x.0)
} else {
None
}
})
.context("in getting the index")?;
let rl = interfaces.len();
if i.name.ends_with(&format!("~{root}")) {
if let ImportKind::Func(f) = i.kind {
let fs = m.funcs[f].sig();
let fname = m.funcs[f].name().to_owned();
let mut b = FunctionBody::new(&m, fs);
let k = b.entry;
let mut values: Vec<waffle::Value> =
b.blocks[k].params.iter().map(|a| a.1).collect();
let mut e = Expr::Bind(
Operator::I32Add,
vec![
Expr::Bind(Operator::I32Const { value: ridx as u32 }, vec![]),
Expr::Bind(
Operator::I32Mul,
vec![
Expr::Bind(Operator::I32Const { value: rl as u32 }, vec![]),
Expr::Leaf(values[0]),
],
),
],
);
let (v, k) = e.build(m, &mut b, k)?;
values[0] = v;
b.set_terminator(k, waffle::Terminator::Return { values });
m.funcs[f] = FuncDecl::Body(fs, fname, b);
continue;
}
} else {
let ek = format!("{}/~{root}{}", i.module, i.name);
let mut ex = m.exports.iter();
let ex = loop {
let Some(x) = ex.next() else {
anyhow::bail!("export not found")
};
if x.name != ek {
continue;
};
let ExportKind::Func(ef) = &x.kind else {
continue;
};
break *ef;
};
if let ImportKind::Func(f) = i.kind {
let fs = m.funcs[f].sig();
let fname = m.funcs[f].name().to_owned();
let mut b = FunctionBody::new(&m, fs);
let k = b.entry;
let mut args = b.blocks[b.entry]
.params
.iter()
.map(|a| a.1)
.collect::<Vec<_>>();
let mut e = Expr::Bind(
Operator::I32DivU,
vec![
Expr::Leaf(args[0]),
Expr::Bind(Operator::I32Const { value: rl as u32 }, vec![]),
],
);
let (v, k) = e.build(m, &mut b, k)?;
args[0] = v;
b.set_terminator(k, waffle::Terminator::ReturnCall { func: ex, args });
m.funcs[f] = FuncDecl::Body(fs, fname, b);
continue;
}
}
}
if i.module == "pit" && i.name == "drop" {
let fs = interfaces
.iter()
.filter_map(|i| {
let mut ex = m.exports.iter();
let ex = loop {
let Some(x) = ex.next() else {
return None;
};
if x.name != format!("pit/{}/{root}.drop", i.rid_str()) {
continue;
};
let ExportKind::Func(ef) = &x.kind else {
continue;
};
break *ef;
};
return Some(ex);
})
.collect::<Vec<_>>();
let t = m.tables.push(TableData {
ty: Type::FuncRef,
initial: fs.len() as u32,
max: Some(fs.len() as u32),
func_elements: Some(fs.clone()),
});
if let ImportKind::Func(f) = i.kind {
let fsi = m.funcs[f].sig();
let fname = m.funcs[f].name().to_owned();
let mut b = FunctionBody::new(&m, fsi);
let k = b.entry;
let mut args = b.blocks[b.entry]
.params
.iter()
.map(|a| a.1)
.collect::<Vec<_>>();
let mut e = Expr::Bind(
Operator::I32DivU,
vec![
Expr::Leaf(b.blocks[k].params[0].1),
Expr::Bind(
Operator::I32Const {
value: fs.len() as u32,
},
vec![],
),
],
);
let (a, k) = e.build(m, &mut b, k)?;
let mut e = Expr::Bind(
Operator::I32RemU,
vec![
Expr::Leaf(b.blocks[k].params[0].1),
Expr::Bind(
Operator::I32Const {
value: fs.len() as u32,
},
vec![],
),
],
);
let (c, k) = e.build(m, &mut b, k)?;
args[0] = a;
args.push(c);
b.set_terminator(
k,
waffle::Terminator::ReturnCallIndirect {
sig: fsi,
table: t,
args,
},
);
m.funcs[f] = FuncDecl::Body(fsi, fname, b);
continue;
}
}
m.imports.push(i);
}
Ok(())
}

129
pit-patch/src/util.rs Normal file
View file

@ -0,0 +1,129 @@
use waffle::{util::new_sig, BlockTarget, Func, FunctionBody, Module, Operator, SignatureData, Table, Type};
use waffle_ast::{add_op, Builder, Expr};
pub fn talloc(m: &mut Module, t: Table) -> anyhow::Result<Func> {
let e = m.tables[t].ty.clone();
let sig = new_sig(
m,
SignatureData {
params: vec![e],
returns: vec![Type::I32],
},
);
let mut f = FunctionBody::new(m, sig);
let n = f.add_block();
let zero = add_op(&mut f, &[], &[Type::I32], Operator::I32Const { value: 0 });
f.append_to_block(f.entry, zero);
f.set_terminator(
f.entry,
waffle::Terminator::Br {
target: BlockTarget {
block: n,
args: vec![zero],
},
},
);
let idx = f.add_blockparam(n, Type::I32);
let mut e = Expr::Bind(
Operator::RefIsNull,
vec![Expr::Bind(
Operator::TableGet { table_index: t },
vec![Expr::Leaf(idx)],
)],
);
let (r, o) = e.build(m, &mut f, n)?;
let mut e = Expr::Bind(
Operator::I32Add,
vec![
Expr::Bind(Operator::I32Const { value: 1 }, vec![]),
Expr::Leaf(idx),
],
);
let (s, o) = e.build(m, &mut f, o)?;
let p = f.add_block();
f.set_terminator(
o,
waffle::Terminator::CondBr {
cond: r,
if_true: BlockTarget {
block: p,
args: vec![],
},
if_false: BlockTarget {
block: n,
args: vec![s],
},
},
);
let q = f.add_block();
let r = f.add_block();
let mut e = Expr::Bind(
Operator::I32Eq,
vec![
Expr::Bind(Operator::TableSize { table_index: t }, vec![]),
Expr::Leaf(idx),
],
);
let (s, o) = e.build(m, &mut f, p)?;
f.set_terminator(
o,
waffle::Terminator::CondBr {
cond: s,
if_true: BlockTarget {
block: r,
args: vec![],
},
if_false: BlockTarget {
block: q,
args: vec![],
},
},
);
let mut e = Expr::Bind(
Operator::TableSet { table_index: t },
vec![Expr::Leaf(idx), Expr::Leaf(f.blocks[f.entry].params[0].1)],
);
let (s, o) = e.build(m, &mut f, q)?;
f.set_terminator(o, waffle::Terminator::Return { values: vec![idx] });
let mut e = Expr::Bind(
Operator::TableGrow { table_index: t },
vec![
Expr::Bind(Operator::I32Const { value: 1 }, vec![]),
Expr::Leaf(f.blocks[f.entry].params[0].1),
],
);
let (s, o) = e.build(m, &mut f, r)?;
f.set_terminator(o, waffle::Terminator::Return { values: vec![idx] });
return Ok(m
.funcs
.push(waffle::FuncDecl::Body(sig, format!("talloc"), f)));
}
pub fn tfree(m: &mut Module, t: Table) -> anyhow::Result<Func> {
let ety = m.tables[t].ty.clone();
let sig = new_sig(
m,
SignatureData {
params: vec![Type::I32],
returns: vec![ety.clone()],
},
);
let mut f = FunctionBody::new(m, sig);
let o = f.entry;
let mut e = Expr::Bind(
Operator::TableGet { table_index: t },
vec![Expr::Leaf(f.blocks[f.entry].params[0].1)],
);
let (r, o) = e.build(m, &mut f, o)?;
let mut e = Expr::Bind(
Operator::TableSet { table_index: t },
vec![
Expr::Leaf(f.blocks[f.entry].params[0].1),
Expr::Bind(Operator::RefNull { ty: ety }, vec![]),
],
);
let (_, o) = e.build(m, &mut f, o)?;
return Ok(m
.funcs
.push(waffle::FuncDecl::Body(sig, format!("tfree"), f)));
}

15
pit-rust-guest/Cargo.toml Normal file
View file

@ -0,0 +1,15 @@
[package]
name = "pit-rust-guest"
version = "0.1.0"
edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
hex = "0.4.3"
pit-core = { version = "0.1.0", path = "../pit-core" }
proc-macro2 = "1.0.85"
quasiquote = "0.1.1"
quote = "1.0.36"
sha3 = "0.10.8"
syn = "2.0.66"

142
pit-rust-guest/src/lib.rs Normal file
View file

@ -0,0 +1,142 @@
use std::iter::once;
use pit_core::{Arg, Interface, ResTy, Sig};
use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use quasiquote::quasiquote;
use sha3::Digest;
use std::io::Write;
pub struct Opts{
pub root: TokenStream,
pub salt: Vec<u8>
}
pub fn render(opts: &Opts, i: &Interface) -> TokenStream{
let root = &opts.root;
let id = i.rid_str();
let mut ha = sha3::Sha3_256::default();
write!(ha,"~{}",id);
ha.write(&opts.salt);
let ha = hex::encode(ha.finalize());
let id2 = format_ident!("R{}",i.rid_str());
let internal = format_ident!("{id}_utils");
let methods = i.methods.iter().map(|(a,b)|quasiquote! {
fn #{format_ident!("{a}")}#{render_sig(root,i,b,&quote! {&self})}
});
let impl_dyns = i.methods.iter().map(|(a,b)|quasiquote! {
fn #{format_ident!("{a}")}#{render_sig(root,i,b,&quote! {&self})}{
#[#root::externref::externref(crate = #{quote! {#root::externref}.to_string()})]
#[wasm_import_module = #{format!("pit/{}",i.rid_str())}]
extern "C"{
#[wasm_import_name = #a]
fn go #{render_sig(root, i,b, &quote! {&#root::externref::Resource<Box<dyn #id2>>})}
};
return go(self,#{
let params = b.params.iter().enumerate().map(|(a,_)|format_ident!("p{a}"));
quote! {
#(#params),*
}
});
}
});
let chains2 = i.methods.iter().map(|(a,b)|quasiquote! {
#[#root::externref::externref(crate = #{quote! {#root::externref}.to_string()})]
#[export_name = #{format!("pit/{id}/~{ha}/{a}")}]
fn #{format_ident!("{a}")}#{render_sig(root,i,b,&quote! {id: *mut Box<dyn #id>})}{
return unsafe{&mut *id}.#{format_ident!("{a}")}(#{
let params = b.params.iter().enumerate().map(|(a,_)|format_ident!("p{a}"));
quote! {
#(#params),*
}
});
}
});
quasiquote!{
pub trait #id2{
#(#methods)*
}
mod #internal{
const _: () = {
#[link_section = ".pit-types"]
static SECTION_CONTENT: [u8; _] = #{
let b = i.to_string();
let c = b.as_bytes().iter().cloned().chain(once(0u8));
quote!{
[#(#c),*]
}
};
};
use super::#id2;
impl #id2 for #root::externref::Resource<Box<dyn #id2>>{
#(#impl_dyns),*
}
#[#root::externref::externref(crate = #{quote! {#root::externref}.to_string()})]
#[export_name = #{format!("pit/{id}/~{ha}.drop")}]
fn _drop(a: *mut Box<dyn #id>){
unsafe{
Box::from_raw(a)
}
}
#(#chains2);*
#[#root::externref::externref(crate = #{quote! {#root::externref}.to_string()})]
#[wasm_import_module = #{format!("pit/{}",i.rid_str())}]
extern "C"{
#[wasm_import_name = #{format!("~{ha}")}]
fn _push(a: *mut Box<dyn #id2>) -> #root::externref::Resource<Box<dyn #id2>>;
}
pub fn push(a: Box<dyn #id2>) -> #root::externref::Resource<Box<dyn #id2>>{
return unsafe{
_push(Box::into_raw(Box::new(a)))
}
}
}
}
}
pub fn render_sig(root: &TokenStream,base: &Interface, s: &Sig, self_: &TokenStream) -> TokenStream{
let params = s.params.iter().map(|a|render_ty(root,base,a)).enumerate().map(|(a,b)|quasiquote!(#{format_ident!("p{a}")} : #b));
let params = once(self_).cloned().chain(params);
let rets = s.rets.iter().map(|a|render_ty(root,base,a));
quote! {
(#(#params),*) -> (#(#rets),*)
}
}
pub fn render_ty(root: &TokenStream,base:&Interface, p: &Arg) -> TokenStream{
match p{
Arg::I32 => quote! {
u32
},
Arg::I64 => quote! {
u64
},
Arg::F32 => quote! {
f32
},
Arg::F64 => quote! {
f64
},
Arg::Resource{ty,nullable,take} => {
let ty = match ty{
ResTy::Of(a) => quasiquote!{
#root::externref::Resource<Box<dyn #{format_ident!("R{}",hex::encode(a))}>>
},
ResTy::None => quote! {
#root::externref::Resource<()>
},
ResTy::This => quasiquote! {
#root::externref::Resource<Box<dyn #{format_ident!("R{}",base.rid_str())}>>
}
};
let ty = if *nullable{
quote! {Option<#ty>}
}else{
ty
};
let ty = if *take{
ty
}else{
quote! {&#ty}
};
ty
},
// Arg::Func(_) => todo!()
}
}

11
pit-rust-host-lib/Cargo.toml Normal file
View file

@ -0,0 +1,11 @@
[package]
name = "pit-rust-host-lib"
version = "0.1.0"
edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
anyhow = "1.0.86"
pit-core = { version = "0.1.0", path = "../pit-core" }
wasm_runtime_layer = "0.4.0"

182
pit-rust-host-lib/src/lib.rs Normal file
View file

@ -0,0 +1,182 @@
use std::{
cell::UnsafeCell,
collections::BTreeMap,
env::vars,
iter::{empty, once},
sync::{Arc, Mutex},
};
use pit_core::{Arg, Interface};
use wasm_runtime_layer::{
backend::WasmEngine, AsContext, AsContextMut, Extern, ExternRef, Func, FuncType, Imports,
Instance, Module, StoreContext, StoreContextMut, Value, ValueType,
};
pub fn init<U: AsRef<Instance> + 'static, E: WasmEngine>(
l: &mut Imports,
ctx: &mut impl AsContextMut<UserState = U, Engine = E>,
){
l.define("pit", "drop", Extern::Func(Func::new(
&mut *ctx,
FuncType::new(
once(ValueType::ExternRef),
empty(),
),
move |mut ctx, args, rets| {
let Value::ExternRef(Some(a)) = args[0].clone() else {
anyhow::bail!("invalid type")
};
let Ok(x): Result<&Wrapped<U,E>, anyhow::Error> = a.downcast::<'_,'_,Wrapped<U,E>,U,E>(ctx.as_context()) else{
return Ok(());
};
let f = x.all[0].clone();
f(ctx.as_context_mut(),vec![])?;
Ok(())
},
)));
}
pub fn emit_ty(a: &Arg) -> ValueType {
match a {
Arg::I32 => ValueType::I32,
Arg::I64 => ValueType::I64,
Arg::F32 => ValueType::F32,
Arg::F64 => ValueType::F64,
Arg::Resource { ty, nullable, take } => ValueType::ExternRef,
// Arg::Func(f) => ValueType::FuncRef,
}
}
pub fn emit<U: AsRef<Instance> + 'static, E: WasmEngine>(
l: &mut Imports,
rid: Arc<Interface>,
m: &Module,
ctx: &mut impl AsContextMut<UserState = U, Engine = E>,
) {
let n = format!("pit/{}", rid.rid_str());
for (j, (i, m)) in rid.methods.iter().enumerate() {
l.define(
&n,
i.as_str(),
Extern::Func(Func::new(
&mut *ctx,
FuncType::new(
once(ValueType::ExternRef).chain(m.params.iter().map(emit_ty)),
m.rets.iter().map(emit_ty),
),
move |mut ctx, args, rets| {
let Value::ExternRef(Some(a)) = args[0].clone() else {
anyhow::bail!("invalid type")
};
let x: &Wrapped<U,E> = a.downcast(ctx.as_context())?;
let t = x.all[j + 1].clone();
let rets2 = t(ctx.as_context_mut(),args[1..].iter().cloned().collect())?;
for (r, s) in rets2.into_iter().zip(rets.iter_mut()) {
*s = r;
}
Ok(())
},
)),
)
}
let i = m.imports(ctx.as_context().engine()).map(|a|(a.module.to_owned(),a.name.to_owned())).collect::<Vec<_>>();
for i in i {
if i.0 == n{
if let Some(t) = i.1.strip_prefix("~"){
let t = t.to_owned();
let rid = rid.clone();
l.define(
&n,
&i.1,
Extern::Func(Func::new(
&mut *ctx,
FuncType::new(once(ValueType::I32), once(ValueType::ExternRef)),
move |mut ctx, args, rets| {
let Value::I32(a) = &args[0] else {
anyhow::bail!("invalid type");
};
let i = ctx.data().as_ref().clone();
let object = Wrapped::new(*a as u32, rid.clone(),t.to_owned(), i, ctx.as_context_mut());
rets[0] = Value::ExternRef(Some(ExternRef::new(ctx, object)));
Ok(())
},
)),
)
};
};
};
}
pub struct Wrapped<U: 'static, E: wasm_runtime_layer::backend::WasmEngine> {
pub rid: Arc<Interface>,
pub all: Vec<Arc<dyn Fn(StoreContextMut<'_,U,E>,Vec<Value>) -> anyhow::Result<Vec<Value>> + Send + Sync>>,
// },
}
impl<U: 'static, E: wasm_runtime_layer::backend::WasmEngine> Wrapped<U,E> {
pub fn new(
base: u32,
rid: Arc<Interface>,
rs: String,
instance: ::wasm_runtime_layer::Instance,
store: ::wasm_runtime_layer::StoreContextMut<'_, U, E>,
) -> Self {
// impl<U: 'static, E: WasmEngine> Copy for X<U, E> {}
// let ctx: X<U, E> = X {
// base: Arc::new(Mutex::new(unsafe { std::mem::transmute(store) })),
// };
let rid2 = rid.clone();
let instance2 = instance.clone();
// let ctx2 = ctx.clone();
let rs2 = rs.clone();
return Self {
rid: rid.clone(),
all: once(
Arc::new(move |mut ctx:StoreContextMut<'_,U,E>,vals: Vec<Value>| -> anyhow::Result<Vec<Value>> {
// let _ = &ctx2;
// let mut b = ctx2.base.lock().unwrap();
let vals: Vec<_> = vec![Value::I32(base as i32)]
.into_iter()
.chain(vals.into_iter())
.collect();
let mut rets = vec![];
let f = instance2.get_export(ctx.as_context_mut(), &format!("pit/{}/~{rs2}.drop", rid2.rid_str()));
let Some(Extern::Func(f)) = f else {
anyhow::bail!("invalid func")
};
f.call(ctx.as_context_mut(), &vals, &mut rets)?;
Ok(rets)
})
as Arc<dyn Fn(StoreContextMut<'_,U,E>,Vec<Value>) -> anyhow::Result<Vec<Value>> + Send + Sync>
)
.chain(rid.methods.iter().map(|(a, b)| {
let rid = rid.clone();
let a = a.clone();
let instance = instance.clone();
let b = Arc::new(b.clone());
// let ctx = ctx.clone();
let rs = rs.clone();
(Arc::new(move |mut ctx:StoreContextMut<'_,U,E>,vals: Vec<Value>| -> anyhow::Result<Vec<Value>> {
// let _ = &ctx;
// let mut bi = ctx.base.lock().unwrap();
let vals: Vec<_> = vec![Value::I32(base as i32)]
.into_iter()
.chain(vals.into_iter())
.collect();
let mut rets = vec![Value::I32(0); b.rets.len()];
let f = instance.get_export(ctx.as_context_mut(), &format!("pit/{}/~{rs}/{a}", rid.rid_str()));
let Some(Extern::Func(f)) = f else {
anyhow::bail!("invalid func")
};
f.call(ctx.as_context_mut(), &vals, &mut rets)?;
Ok(rets)
})
as Arc<dyn Fn(StoreContextMut<'_,U,E>,Vec<Value>) -> anyhow::Result<Vec<Value>> + Send + Sync>)
}))
.collect(),
};
}
}
// impl Drop for Wrapped {
// fn drop(&mut self) {
// self.all[0](vec![]).unwrap();
// }
// }
pub type RWrapped<U,E> = ::std::sync::Arc<Wrapped<U,E>>;
pub extern crate anyhow;
pub extern crate wasm_runtime_layer;

14
pit-rust-host/Cargo.toml Normal file
View file

@ -0,0 +1,14 @@
[package]
name = "pit-rust-host"
version = "0.1.0"
edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
hex = "0.4.3"
pit-core = { version = "0.1.0", path = "../pit-core" }
proc-macro2 = "1.0.85"
quasiquote = "0.1.1"
quote = "1.0.36"
syn = "2.0.66"

291
pit-rust-host/src/lib.rs Normal file
View file

@ -0,0 +1,291 @@
use std::iter::once;
use pit_core::{Arg, Interface, Sig};
use proc_macro2::TokenStream;
use quasiquote::quasiquote;
use quote::{format_ident, quote, ToTokens};
use syn::{spanned::Spanned, Ident, Index};
pub fn render(root: &TokenStream, i: &Interface) -> TokenStream {
let id = format_ident!("R{}", i.rid_str());
let internal = format_ident!("{id}_utils");
let methods = i.methods.iter().map(|(a, b)| {
quasiquote! {
fn #{format_ident!("{a}")}#{render_sig(root,b,&quote! {&self})}
}
});
let impls = i.methods.iter().map(|(a, b)| {
let init = b.params.iter().enumerate().map(|(pi,a)|render_new_val(root, a, quasiquote!{#{format_ident!("p{pi}")}}));
let init = once(quote!{#root::wasm_runtime_layer::Value::I32(self.base as i32)}).chain(init);
let fini = b.rets.iter().enumerate().map(|(ri,r)|{
quasiquote!{
#{render_base_val(root, r, quote! {&rets[#ri]})}
}
});
quasiquote! {
fn #{format_ident!("{a}")}#{render_sig(root,b,&quote! {&self})}{
let Some(#root::wasm_runtime_layer::Export::Func(f)) = self.instance.get_export(unsafe{
&*::std::cell::UnsafeCell::raw_get(self.ctx)
},#{format!("pit/{}/{}",i.rid_str(),a)}) else{
panic!("invalid func")
};
let args = vec![#(#init),*];
let mut rets = vec![#root::wasm_runtime_layer::Value::I32(0);#{b.rets.len()}];
f.call(unsafe{
&mut *::std::cell::UnsafeCell::raw_get(self.ctx)
},&args,&mut rets)?;
return Ok((#(#fini),*))
}
}
});
let injects = i.methods.iter().map(|(a,b)|{
let init = b.params.iter().enumerate().map(|(pi,a)|quasiquote!{let #{format_ident!("p{pi}")} = #{render_base_val(root, a, quote! {&args[#pi + 1]})}});
let fini = b.rets.iter().enumerate().map(|(ri,r)|{
quasiquote!{
rets[#ri] = #{render_new_val(root, r, quasiquote! {r . #{Index{index: ri as u32,span: root.span()}}})}
}
});
quasiquote!{
let f = #root::wasm_runtime_layer::Func::new(ctx,#{render_blit_sig(root, b)},|ctx,args,rets|{
#(#init),*
let #root::wasm_runtime_layer::Value::ExternRef(t) = &args[0] else{
#root::anyhow::bail!("invalid param")
}
let r = match t.downcast::<'_,'_, ::std::sync::Arc<dyn #id<U,E>>,U,E>(ctx){
Ok(t) => {
t.#{format_ident!("{a}")}(#{
let p = b.params.iter().enumerate().map(|(a,_)|format_ident!("p{a}"));
quote! {
#(#p .clone()),*
}
})?
},
Err(_) => match t.downcast::<'_,'_, ::std::sync::Arc<#root::Wrapped<U,E>>,U,E>(ctx)?{
Ok(t) => {
t.#{format_ident!("{a}")}(#{
let p = b.params.iter().enumerate().map(|(a,_)|format_ident!("p{a}"));
quote! {
#(#p .clone()),*
}
})?
},
Err(_) => #root::anyhow::bail!("invalid externref")
}
};
#(#fini),*
Ok(())
});
i.define(#{format!("pit/{}",i.rid_str())},#a,f);
}}).chain(once( quasiquote!{
let f = #root::wasm_runtime_layer::Func::new(ctx,#root::wasm_runtime_layer::FuncType::new([#root::wasm_runtime_layer::ValueType::I32],[#root::wasm_runtime_layer::ExternType]),|ctx,args,rets|{
let #root::wasm_runtime_layer::Value::I32(i) = &args[0] else{
unreachable!()
};
let n = ctx.data().as_ref().clone();
let v = #root::Wrapped::new(*i as u32,#{i.rid_str()},ctx.as_context_mut(),n);
// let v: ::std::sync::Arc<dyn #id<U,E>> = ::std::sync::Arc::new(v);
rets[0] = #root::wasm_runtime_layer::Value::ExternRef(#root::wasm_runtime_layer::ExternRef::new(ctx,v));
Ok(())
});
i.define(#{format!("pit/{}",i.rid_str())},"~push",f);
}));
quasiquote! {
pub trait #id<U: 'static,E: #root::wasm_runtime_layer::backend::WasmEngine>{
fn to_any(&self) -> &dyn ::std::any::Any;
#(#methods)*
}
mod #internal{
use super::#id;
impl<U: 'static,E: #root::wasm_runtime_layer::backend::WasmEngine> #id<U,E> for #root::RWrapped<U,E>{
fn to_any(&self) -> &dyn ::std::any::Any{
return self;
}
#(#impls)*
}
fn inject<U: 'static + AsRef<#root::wasm_runtime_layer::Instance>,E: #root::wasm_runtime_layer::backend::WasmEngine>(ctx: impl #root::wasm_runtime_layer::AsContextMut<UserState = U,Engine = E>, i: &mut #root::wasm_runtime_layer::Imports){
#(#injects)*
}
}
}
}
pub fn render_sig(root: &TokenStream, s: &Sig, self_: &TokenStream) -> TokenStream {
let params = s
.params
.iter()
.map(|a| render_ty(root, a))
.enumerate()
.map(|(a, b)| quasiquote!(#{format_ident!("p{a}")} : #b));
let params = once(self_).cloned().chain(params);
let rets = s.rets.iter().map(|a| render_ty(root, a));
quote! {
(#(#params),*) -> #root::anyhow::Result<(#(#rets),*)>
}
}
pub fn render_blit(root: &TokenStream, p: &Arg) -> TokenStream {
match p {
Arg::I32 => quote! {
#root::wasm_runtime_layer::ValueType::I32
},
Arg::I64 => quote! {
#root::wasm_runtime_layer::ValueType::I64
},
Arg::F32 => quote! {
#root::wasm_runtime_layer::ValueType::F32
},
Arg::F64 => quote! {
#root::wasm_runtime_layer::ValueType::F64
},
Arg::Resource { .. } => quote! {
#root::wasm_runtime_layer::ValueType::ExternRef
},
}
}
pub fn render_blit_sig(root: &TokenStream, s: &Sig) -> TokenStream {
quasiquote! {
#root::wasm_runtime_layer::FuncType::new([#{
let p = s.params.iter().map(|p|render_blit(root, p));
let p = once(quote! {#root::wasm_runtime_layer::ValueType::ExternRef}).chain(p);
quote! {
#(#p),*
}
}].into_iter(),[#{ let p = s.rets.iter().map(|p|render_blit(root, p));
quote! {
#(#p),*
}}].into_iter())
}
}
pub fn render_ty(root: &TokenStream, p: &Arg) -> TokenStream {
match p {
Arg::I32 => quote! {
u32
},
Arg::I64 => quote! {
u64
},
Arg::F32 => quote! {
f32
},
Arg::F64 => quote! {
f64
},
Arg::Resource { ty, nullable } => match ty {
None => quote! {
#root::wasm_runtime_layer::ExternRef
},
Some(a) => {
let a = quasiquote! {
::std::sync::Arc<dyn #{format_ident!("R{}",hex::encode(a))}<U,E>>
};
if *nullable {
quote! {Option<#a>}
} else {
a
}
}
},
}
}
pub fn render_base_val(root: &TokenStream, p: &Arg, x: TokenStream) -> TokenStream {
let v = match p {
Arg::I32 => quote! {
let #root::wasm_runtime_layer::Value::I32(t) = #x else{
#root::anyhow::bail!("invalid param")
}
},
Arg::I64 => quote! {
let #root::wasm_runtime_layer::Value::I64(t) = #x else{
#root::anyhow::bail!("invalid param")
}
},
Arg::F32 => quote! {
let #root::wasm_runtime_layer::Value::F32(t) = #x else{
#root::anyhow::bail!("invalid param")
}
},
Arg::F64 => quote! {
let #root::wasm_runtime_layer::Value::F64(t) = #x else{
#root::anyhow::bail!("invalid param")
}
},
Arg::Resource { ty, nullable } => {
let mut a = quote! {
let #root::wasm_runtime_layer::Value::ExternRef(t) = #x else{
#root::anyhow::bail!("invalid param")
}
};
if let Some(r) = ty.as_ref() {
quasiquote!{
let t = match t{None => None,Some(t) => Some(match t.downcast::<'_,'_, ::std::sync::Arc<dyn #{format_ident!("R{}",hex::encode(r))}<U,E>>,U,E>(ctx){
Ok(t) => t.clone(),
Err(_) => match t.downcast::<'_,'_, ::std::sync::Arc<#root::Wrapped<U,E>>,U,E>(ctx){
Ok(t) => Arc::new(t.clone()),
Err(_) => #root::anyhow::bail!("invalid param")
}
})}
}.to_tokens(&mut a);
}
if !*nullable {
quote! {
let t = match t{
Some(a) => a,
None => #root::anyhow::bail!("invalid param")
}
}
.to_tokens(&mut a)
}
a
}
};
quote! {
{
#v;t
}
}
}
pub fn render_new_val(root: &TokenStream, p: &Arg, t: TokenStream) -> TokenStream {
match p {
Arg::I32 => quote! {
#root::wasm_runtime_layer::Value::I32(#t)
},
Arg::I64 => quote! {
#root::wasm_runtime_layer::Value::I64(#t)
},
Arg::F32 => quote! {
#root::wasm_runtime_layer::Value::F32(#t)
},
Arg::F64 => quote! {
#root::wasm_runtime_layer::Value::F64(#t)
},
Arg::Resource { ty, nullable } => {
let tq = |t:TokenStream|quasiquote! {
{
let t = #t;
#{match ty{
None => quote! { #root::wasm_runtime_layer::ExternRef::new(ctx,t)},
Some(_) => quote! {
match t.to_any().downcast_ref::<::std::sync::Arc<#root::Wrapped<U,E>>>(){
None => #root::wasm_runtime_layer::ExternRef::new(ctx,t),
Some(t) => #root::wasm_runtime_layer::ExternRef::new(ctx,t.clone()),
}
}
}}
}
};
if !*nullable {
quasiquote! {
#root::wasm_runtime_layer::Value::ExternRef(Some(#{match ty{
None => t,
Some(_) => tq(t)
}}))
}
} else {
quasiquote! {
#root::wasm_runtime_layer::Value::ExternRef(#{match ty{
None => t,
Some(_) => quasiquote! {#t.map(|t|#{tq(quote! {t})})}
}})
}
}
}
}
}