├── .gitignore ├── Cargo.toml ├── README.md ├── lineiform ├── Cargo.toml └── src │ ├── block.rs │ ├── closure.rs │ ├── lib.rs │ ├── lift.rs │ ├── marker.rs │ └── tracer.rs ├── src ├── main.rs └── parser.rs └── tangle ├── Cargo.toml └── src ├── abi.rs ├── flags.rs ├── ir.rs ├── lib.rs ├── node.rs ├── oper.rs ├── opt.rs ├── port.rs ├── region.rs ├── select.rs └── time.rs /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | .gdb_history 3 | words 4 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "simplejit" 3 | version = "0.1.0" 4 | authors = ["Charlie Cummings "] 5 | edition = "2018" 6 | 7 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 8 | 9 | [dependencies] 10 | thiserror = "1.0" 11 | enum-display-derive = "0.1.1" 12 | yaxpeax-arch = "0.2.5" 13 | goblin = "0.4" 14 | target-lexicon = "0.12.1" 15 | bitvec = "0.22.3" 16 | bitflags = "1.3.1" 17 | rangemap = "0.1.11" 18 | lineiform = { path = "./lineiform" } 19 | frunk = "0.4.0" 20 | 21 | [dependencies.yaxpeax-x86] 22 | version = "1.0.4" 23 | features = ["fmt"] 24 | path = "../yaxpeax-x86" # FIXME when yax x86 1.1 comes out 25 | 26 | [profile.dev] 27 | opt-level = 2 # we use tailcalls a lot, which will blow the stack on O0 28 | 29 | [profile.release] 30 | debug = false 31 | opt-level = 3 32 | overflow-checks = true 33 | panic = "abort" 34 | 35 | [workspace] 36 | 37 | members = [ 38 | "tangle", 39 | "lineiform" 40 | ] 41 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Lineiform 2 | 3 | Lineiform is a meta-JIT library for Rust interpreters. Given an interpreter that uses closure generation, it allows an author to add minimal annotations and calls into Lineiform in order to automagically get an optimizing method JIT, via runtime function inlining and constant propagation to resolve dynamic calls to closed environment members. Internally it does lifting of closure bodies from x86 assembly to Cranelift IR for dynamic recompilation. 4 | 5 | # The Big Idea 6 | (There's a [blog post](https://blog.redvice.org/2022/lineiform-rust-meta-jit/) now). 7 | 8 | A normal tree-walking interpreter in Rust looks like this: (pseudocode but not much) 9 | ```rust 10 | fn eval(ast: Ast) -> usize { 11 | match ast { 12 | Literal(l) => l, 13 | Add(left, right) => eval(left) + eval(right), 14 | } 15 | } 16 | 17 | let val = eval(Add(Literal(1), Literal(2))); 18 | val // 3 19 | ``` 20 | 21 | You can instead write a *closure generator* interpreter like this, for the same language: (also pseudocode but not much) 22 | 23 | ```rust 24 | type Environment = (); 25 | type EvalAst = Fn(Environment) -> usize; 26 | fn eval(ast: Ast) -> EvalAst { 27 | match ast { 28 | Literal(l) => |e: Environment| l, 29 | Add(_left, _right) => { 30 | let left = eval(_left); 31 | let right = eval(_right); 32 | |e: Environment| { left(e) + right(e) } 33 | } 34 | } 35 | 36 | let compiled_eval = eval(Add(Literal(1), Literal(2))) 37 | let env: Environment = (); 38 | compiled_eval(env) // 3 39 | ``` 40 | 41 | `compiled_eval` look like `[ return [ return 1] + [ return 2] ]` where `[]` is its own closure. 42 | The Add case of the closure in particular has a closed environment of `left` and `right` which are a specialized closure for the AST node that they're made of - it's just that you still need a dynamic call for evaluating them. Lineiform is a library that you'd call (probably only in a `Call` language construct, for method-level optimization) to optimize an EvalAst closure, for example inlining those `left` and `right` dynamic calls into the single function body of `compiled_eval`, since the closed environment for an EvalAst is statically known at that point. 43 | `compiled_eval` then become `[ return 1 + 2 ]`, which you can then optimize further to just `[ return 3]` - it "unrolled" multiple interpreter loops into one function, exposing them to a recompiler for global optimization of the actual evaluation of the closure. We don't fix the `Environment` that is actually *passed* to the closure, only the closed environment. 44 | 45 | You can think of this analogous to how a Forth compiler works: `1` is a word that pushes the value 1, `2` is a word that pushes the value 2. A dumb value of `1 2 +` then is three dynamic function calls, to function body addresses stored in its "contents" list (`1`, `2`, `+`). If you have a concretely built word at runtime, however, you can inline the calls to all the function pointers you have, and get the *actual* execution path of the function, allowing you to do global optimizations. The Rust closure `compiled_eval` looks much like this Forth word if you squint, but doesn't use a threaded interpreter execution mode, and stores the function pointers in the closure closed environment instead of a word body array. 46 | 47 | # Lifting x86 48 | Lineiform uses a technique called dynamic recompilation, which is used extensively in cybersecurity libraries like DynamoRIO and emulators like QEMU: by re-encoding the semantics of x86 assembly as some internal IR ("lifting"), you can treat the IR much like the IR you'd build from a normal compiler frontend. This allows you to run normal optimization passes over the IR and compile it back down to assembly again ("lowering"). 49 | 50 | Lineiform uses Cranelift's IR as its lifting target, allowing us to use its existing infrastructure to reallocate registers after our function inlining and JIT compile back to x86. This isn't really what it was built for, tbh, so it doesn't do all that great of a job currently. 51 | 52 | ## What works 53 | You can create a JIT and feed it a closure, which does some constant propagation from the closure's environment in order to turn tailcalls and dynamic function calls into inlined basic blocks (and inline any constants that were closed over, for example) and call the resulting "optimized" closure. The x86 lifting isn't very featureful - I've implemented register aliasing, memory load/store widths, condition flags and some branching - but I haven't done any verification that any operations are correct w.r.t. actual instruction execution. 54 | 55 | ## What doesn't 56 | The main missing lifting features are 1) it never merges divergant flow control back together 2) it doesn't handle loops at all 3) all operations are (currently) zero extended, instead of being zero or sign extended depending on the instruction. And just a whole bunch of instructions, of course, since I've mostly only been implementing exactly as much as I require to pass each test I add. C'est la vie. 57 | 58 | # Warning 59 | This is extremely unsafe and under construction. It will silently miscompile closures so that they return the wrong values and/or eat your socks. The entire idea of using dynamic recompilation for a JIT is ill informed and foolish; the author is the epitome of hubris. There are dragons contained within. 60 | 61 | No license or warranty is provided. If you want to use it and care about licensing things, get in touch [on Twitter](https://twitter.com/sickeningsprawl). 62 | 63 | # TODO list 64 | - [x] Disassembling functions 65 | - [x] Lifting x86 to Cranelift IR 66 | - [ ] Enough instructions that anything non-trivial doesn't crash 67 | - [ ] Tear out Cranelift and use our home-rolled Tangle IR instead 68 | - [x] Branches 69 | - [ ] Merging divergant flow back together 70 | - [ ] Loops 71 | - [x] Calls 72 | - [ ] External calls (we need to "concretize" the SSA value map, so that an external call has the correct registers/stack, and restore them afterwards) 73 | - [ ] Lifting bailouts: if we hit an instruction we can't lift, we can instead bail out and emit a call to inside the original version of the function, so we can optimize only half of a complex function for example. 74 | - [ ] Optimistic external calls: try to use DWARF to figure out calling convention + argument count for external functions, to avoid having to spill more registers/stack than we need! 75 | - [ ] Correct input/output function prototype detection: currently we just hardcode parameter/return value types, which is Bad. 76 | - [ ] Freezing allocator scheme 77 | - [ ] *Freeing* from the frozen allocator (probably using crossbeam_epoch) 78 | - [ ] Smarter function size fetching (currently we only search `/proc/self/exe` for them! Which doesn't work if you have any dynamic linking lol) 79 | - [ ] Performance counters for the "optimized" Fn trait, so we only optimize over some threshold. 80 | - [ ] A `specialize` function: we can have a variant of `optimize` that also fixes a value from the *arguments* to the closure, to for example pin the type tag of an object to erase type checks. We can compile a variant for each concrete pinned value, and emit a jump table that dispatches to either the non-compiled generic version, or the specialized variant. 81 | -------------------------------------------------------------------------------- /lineiform/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "lineiform" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | nom = "7.0.0-alpha1" 10 | jemallocator = "0.3.2" 11 | thiserror = "1.0" 12 | enum-display-derive = "0.1.1" 13 | goblin = "0.4" 14 | lazy_static = "1.4.0" 15 | bitvec = "0.22.3" 16 | bitflags = "1.3.1" 17 | rangemap = "0.1.11" 18 | tangle = { path = "../tangle" } 19 | frunk = "0.4.0" 20 | 21 | yaxpeax-arch = "0.2.5" 22 | 23 | [dependencies.yaxpeax-x86] 24 | version = "1.0.4" 25 | features = ["fmt"] 26 | path = "../../yaxpeax-x86" # FIXME when yax x86 1.1 comes out 27 | 28 | [lib] 29 | path = "src/lib.rs" 30 | -------------------------------------------------------------------------------- /lineiform/src/block.rs: -------------------------------------------------------------------------------- 1 | use crate::tracer::{self, Tracer}; 2 | use std::sync::Arc; 3 | use yaxpeax_x86::long_mode::Instruction; 4 | use yaxpeax_x86::long_mode::RegSpec; 5 | use crate::lift::JitValue; 6 | use std::num::Wrapping; 7 | 8 | use std::ffi::c_void; 9 | use std::marker::PhantomData; 10 | use crate::lift::Location; 11 | pub struct Function { 12 | pub base: *const (), 13 | /// Our offset from function base to wanted IP. (bytes, instructions) 14 | pub offset: (usize, usize), 15 | pub size: usize, 16 | pub instructions: Arc>, 17 | pub pinned: Arc>, 18 | _phantom: PhantomData<(ARG, OUT)>, //fn(data: *const c_void, ARG)->OUT 19 | } 20 | 21 | use thiserror::Error; 22 | #[derive(Error, Debug)] 23 | pub enum BlockError { 24 | #[error("decompilation error")] 25 | Decompile(#[from] tracer::TracerError), 26 | } 27 | 28 | ///! When we lift an Fn(A)->O closure, we actually lift c_fn::, which takes 29 | ///the Fn::call variant for the closure, and a *const c_void closure object, along 30 | ///with a tuple of actual arguments. 31 | ///This is an extern C function, because we need a stable ABI - but making it 32 | ///extern C means that when we *lower* the function back down, we are emitting 33 | ///the extern C prototype function to call! Because if you are using this as 34 | ///a JIT for a closure generator compiler you need to return a *closure* for your 35 | ///function. E.g. `Lineiform::jit(box move |a: A| -> O { o }) -> Fn(A, Output=O)` 36 | ///Which means we need to have a Rust->SystemV->Rust ABI dance: we provide a 37 | ///rust-call interface, which calls our extern C tracing trampoline/lift result, 38 | ///which calls whatever Rust closure body we inlined. We can at least hope that 39 | ///the last inlining reduces register waste. 40 | ///Maybe in the future Rust will allow for custom calling convention of closures, 41 | ///but the issue has ~0 traffic and no RFC, so I'm not holding my breath... 42 | impl Function { 43 | pub fn assume_args(mut self, pinning: Vec<(Location, JitValue)>) -> Self { 44 | self.pinned = Arc::new(pinning); 45 | self 46 | } 47 | 48 | pub fn new(tracer: &mut Tracer, f: *const ()) -> Result { 49 | let base = tracer.get_base()?; 50 | let f_sym = tracer.get_symbol_from_address(f)?; 51 | // Our address could be in the middle of a function - we want to disassemble 52 | // from the start so that we can cache disassembly, and so that when we 53 | // are emitting from it we have the entire function for backwards jmp analysis 54 | let f_base = (base + (f_sym.st_value as usize)); 55 | let f_off = (f as usize) - f_base; 56 | let instructions = tracer.disassemble(f_base as *const (), f_sym.st_size as usize)?; 57 | let mut inst_off = 0; 58 | let mut inst_addr = f_base; 59 | // unfortunately, due to x86 having irregular instruction widths, we can't 60 | // easily turn bytes offset -> instruction stream offset. 61 | // maybe we can put this in a cache as well? or at least have some skiplist 62 | // like structure. 63 | use crate::yaxpeax_arch::LengthedInstruction; 64 | for inst in instructions.iter() { 65 | if inst_addr == (f as usize) { 66 | break; 67 | } else if inst_addr > (f as usize) { 68 | panic!("trying to create a function for {:x} overshot! is it unaligned inside {:x}?", f as usize, f_base); 69 | } else { 70 | inst_off += 1; 71 | inst_addr += inst.len().to_const() as usize; 72 | } 73 | } 74 | Ok(Self { 75 | base: f_base as *const (), 76 | offset: (f_off, inst_off), 77 | size: f_sym.st_size as usize, 78 | instructions: instructions, 79 | pinned: Arc::new(Vec::new()), 80 | _phantom: PhantomData, 81 | }) 82 | } 83 | } 84 | -------------------------------------------------------------------------------- /lineiform/src/closure.rs: -------------------------------------------------------------------------------- 1 | use thiserror::Error; 2 | use std::collections::HashMap; 3 | use crate::parser::{self, Atom, BuiltIn, Expr, parse_expr, ListExpr}; 4 | use std::error::Error as StdError; 5 | use std::num::Wrapping; 6 | 7 | impl Atom { 8 | fn atom(u: usize) -> Expr { 9 | Expr::Constant(Atom::Num(Wrapping(u))) 10 | } 11 | 12 | fn atom_wrap(u: parser::Int) -> Expr { 13 | Expr::Constant(Atom::Num(u)) 14 | } 15 | } 16 | 17 | #[derive(thiserror::Error, Debug)] 18 | pub enum EvalError { 19 | #[error("the variable {0} would shadow")] 20 | Shadow(String), 21 | #[error("the variable {0} is unbound at the unsite")] 22 | Unbound(String), 23 | #[error("wanted type `{0}`, got `{1}`")] 24 | TypeError(&'static str, &'static str), 25 | #[error("runtime panic {0}")] 26 | Panic(&'static str), 27 | #[error("compiling dynamic code failed: {0}")] 28 | DynamicFail(#[from] CompileError), 29 | } 30 | 31 | pub struct Env { 32 | pub variables: HashMap, 33 | } 34 | 35 | impl Env { 36 | fn new() -> Self { 37 | Self { 38 | variables: HashMap::new() 39 | } 40 | } 41 | } 42 | 43 | pub trait ClosureFn = Fn(&mut Env) -> Result; 44 | type ClosureFnPtr = *const (); // fn(&Self, &mut Env<'a>) -> Result 45 | pub(crate) type ClosureExpr = Box< 46 | dyn ClosureFn + Send + Sync, 47 | >; 48 | 49 | use core::ptr::*; 50 | use core::ptr::DynMetadata; 51 | use core::mem; 52 | use core::ptr; 53 | pub fn get_ptr_from_closure(f: &T) -> usize { 54 | let fn_call = ::call_once; 55 | let (addr, meta) = (&*f as &dyn ClosureFn as *const dyn ClosureFn).to_raw_parts(); 56 | #[derive(Debug)] 57 | #[repr(C)] 58 | struct RawMeta { 59 | vtable: &'static [usize;4] 60 | } 61 | unsafe { 62 | let vtable = mem::transmute::<_, RawMeta>(addr); 63 | println!("meta = {:x?}", vtable.vtable[1]); 64 | let methods = mem::transmute::<_, RawMeta>(vtable.vtable[1]); 65 | println!("virtual functions: {:x?}", methods); 66 | let call = methods.vtable[3]; 67 | println!("closure body: {:x}", call); 68 | } 69 | let fn_call_as_ptr = fn_call as usize; 70 | println!("fn_call addr = 0x{:x}", fn_call_as_ptr); 71 | fn_call_as_ptr 72 | } 73 | 74 | #[derive(thiserror::Error, Debug)] 75 | pub enum CompileError { 76 | #[error("nom parsing error {0}")] 77 | Parse(String), 78 | #[error("too few arguments to {0}: wanted {1}, got {2}")] 79 | TooFew(BuiltIn, usize, usize), 80 | #[error("unknown data store error")] 81 | Unknown, 82 | #[error("couldn't compile quote expr {0}")] 83 | QuoteCompile(ListExpr), 84 | } 85 | pub fn compile_expr(src: &str) -> Result { 86 | parse_expr(src) 87 | .map_err(|e| CompileError::Parse(format!("{}", e))) 88 | .and_then(|(_, exp)| generate_closure(exp)) 89 | } 90 | 91 | pub fn generate_closure(expr: Expr) -> Result { 92 | println!("application on {:?}", expr.clone()); 93 | match expr { 94 | Expr::Constant(_) | Expr::Quote(_) => Ok(box move |e|{ Ok(expr.clone()) }), 95 | Expr::IfElse(cond, t, f) => { 96 | let cond = generate_closure(*cond)?; 97 | let t = generate_closure(*t)?; 98 | let f = generate_closure(*f)?; 99 | Ok(box move |e| { 100 | if cond(e)?.need_int()? == Wrapping(0) { 101 | f(e) 102 | } else { 103 | t(e) 104 | } 105 | }) 106 | }, 107 | Expr::Application(op, mut args) => { 108 | match op.clone() { 109 | box Expr::Constant(Atom::BuiltIn(opcode)) => match opcode { 110 | BuiltIn::Equal => { 111 | if let [x, y] = args.as_slice() { 112 | let x = generate_closure(x.clone())?; 113 | let y = generate_closure(y.clone())?; 114 | Ok(box move |e| { 115 | Ok(Atom::atom( 116 | if x(e)?.need_int()? == y(e)?.need_int()? { 117 | 1 118 | } else { 119 | 0 120 | } 121 | )) 122 | }) 123 | } else { 124 | Err(CompileError::TooFew(opcode, 2, args.len())) 125 | } 126 | }, 127 | BuiltIn::Plus => { 128 | if let [x, y] = args.as_slice() { 129 | let x = generate_closure(x.clone())?; 130 | let y = generate_closure(y.clone())?; 131 | Ok(box move |e| { 132 | Ok(Atom::atom_wrap( 133 | x(e)?.need_int()? + y(e)?.need_int()? 134 | )) 135 | }) 136 | } else { 137 | Err(CompileError::TooFew(opcode, 2, args.len())) 138 | } 139 | }, 140 | BuiltIn::Times => { 141 | if let [x, y] = args.as_slice() { 142 | let x = generate_closure(x.clone())?; 143 | let y = generate_closure(y.clone())?; 144 | Ok(box move |e| { 145 | Ok(Atom::atom_wrap( 146 | x(e)?.need_int()? * y(e)?.need_int()? 147 | )) 148 | }) 149 | } else { 150 | Err(CompileError::TooFew(opcode, 2, args.len())) 151 | } 152 | }, 153 | BuiltIn::Divide => { 154 | if let [x, y] = args.as_slice() { 155 | let x = generate_closure(x.clone())?; 156 | let y = generate_closure(y.clone())?; 157 | Ok(box move |e| { 158 | Ok(Atom::atom_wrap( 159 | x(e)?.need_int()? / y(e)?.need_int()? 160 | )) 161 | }) 162 | } else { 163 | Err(CompileError::TooFew(opcode, 2, args.len())) 164 | } 165 | }, 166 | 167 | BuiltIn::Let => { 168 | if let [name, val, cont] = args.as_slice() { 169 | let name = generate_closure(name.clone())?; 170 | let val = generate_closure(val.clone())?; 171 | let cont = generate_closure(cont.clone())?; 172 | println!("build let"); 173 | Ok(box move |e| { 174 | let conc_name = &*name(e)?.need_keyword()?; 175 | let conc_value = val(e)?; 176 | e.variables.try_insert(conc_name.to_string(), conc_value).map_err(|e| EvalError::Shadow(conc_name.to_string()))?; 177 | let r = cont(e); 178 | e.variables.remove(conc_name); 179 | r 180 | }) 181 | } else { 182 | Err(CompileError::TooFew(opcode, 3, args.len())) 183 | } 184 | }, 185 | BuiltIn::Set => { 186 | if let [name, val] = args.as_slice() { 187 | let name = generate_closure(name.clone())?; 188 | let val = generate_closure(val.clone())?; 189 | Ok(box move |e| { 190 | let conc_name = &*name(e)?.need_keyword()?; 191 | let conc_value = val(e)?; 192 | let entry = e.variables.entry(conc_name.to_string()); 193 | if let std::collections::hash_map::Entry::Vacant(_) = entry { 194 | Err(EvalError::Unbound(conc_name.to_string())) 195 | } else { 196 | entry.and_modify(|v| *v = conc_value.clone()); 197 | Ok(conc_value) 198 | } 199 | }) 200 | } else { 201 | Err(CompileError::TooFew(opcode, 3, args.len())) 202 | } 203 | }, 204 | BuiltIn::Get => { 205 | if let [name] = args.as_slice() { 206 | let name = generate_closure(name.clone())?; 207 | Ok(box move |e| { 208 | let conc_name = &*name(e)?.need_keyword()?; 209 | e.variables.get(&conc_name.to_string()) 210 | .map(|v| Ok(v.clone())) 211 | .unwrap_or_else(|| Err(EvalError::Unbound(conc_name.to_string()))) 212 | }) 213 | } else { 214 | Err(CompileError::TooFew(opcode, 1, args.len())) 215 | } 216 | }, 217 | BuiltIn::Do => { 218 | let mut comp_args = args.drain(..).map(|v| generate_closure(v) ) 219 | .collect::, _>>()?; 220 | Ok(box move |e| { 221 | Ok(comp_args.iter().try_fold(Expr::Constant(Atom::Unit), |i, v| v(e))?) 222 | }) 223 | }, 224 | BuiltIn::Loop => { 225 | if let [times, body] = args.as_slice() { 226 | let times = generate_closure(times.clone())?; 227 | let body = generate_closure(body.clone())?; 228 | Ok(box move |e| { 229 | let times = times(e)?.need_int()?; 230 | let mut v = Expr::Constant(Atom::Unit); 231 | for i in 0..times.0 { 232 | v = body(e)?; 233 | } 234 | Ok(v) 235 | }) 236 | } else { 237 | Err(CompileError::TooFew(opcode, 2, args.len())) 238 | } 239 | }, 240 | _ => unimplemented!("unimplemented opcode {}", opcode) 241 | }, 242 | box Expr::Quote(ref v) => { 243 | let v = v.clone(); 244 | Ok(box move |e| { 245 | if let [head, body @ ..] = v.as_slice() { 246 | generate_closure(Expr::Application(box head.clone(), body.to_vec()))?(e) 247 | } else { 248 | Err(EvalError::DynamicFail(CompileError::QuoteCompile(ListExpr(v.to_vec())))) 249 | } 250 | }) 251 | }, 252 | box dynamic @ _ => { 253 | println!("??? {:?}", dynamic); 254 | // We weren't able to build the closure from a constant expression 255 | let op = generate_closure(*op.clone())?; 256 | let args = args.clone().iter() 257 | .map(|elem| generate_closure(elem.clone())) 258 | .collect::, CompileError>>()?; 259 | Ok(box move |e| { 260 | let args = args.iter().map(|elem| elem(e)) 261 | .collect::, EvalError>>()?; 262 | generate_closure(Expr::Application(box op(e)?, args))?(e) 263 | }) 264 | }, 265 | e @ _ => unimplemented!("unimplemented application type {:?}", e) 266 | } 267 | }, 268 | _ => unimplemented!("unimplemented: {:?}", expr.clone()), 269 | } 270 | } 271 | 272 | #[cfg(test)] 273 | mod test { 274 | use crate::*; 275 | use crate::parser::Atom; 276 | use crate::closure::Env; 277 | #[derive(Error, Debug)] 278 | enum TestError { 279 | #[error("failed to compile: {0}")] 280 | Compile(#[from] CompileError), 281 | #[error("failed to evaluate: {0}")] 282 | Eval(#[from] EvalError), 283 | } 284 | 285 | #[test] 286 | fn test_plus() -> Result<(), TestError> { 287 | assert_eq!(compile_expr("(+ 1 2)")?(&mut Env::new())?, Atom::atom(3)); 288 | Ok(()) 289 | } 290 | 291 | #[test] 292 | fn test_plus_eval() -> Result<(), TestError> { 293 | assert_eq!(compile_expr("(+ 1 (+ 2 3))")?(&mut Env::new())?, Atom::atom(6)); 294 | Ok(()) 295 | } 296 | 297 | #[test] 298 | fn test_let() -> Result<(), TestError> { 299 | assert_eq!(compile_expr("(let :x 1 2)")?(&mut Env::new())?, Atom::atom(2)); 300 | Ok(()) 301 | } 302 | 303 | #[test] 304 | fn test_get() -> Result<(), TestError> { 305 | assert_eq!(compile_expr("(let :x 1 (get :x))")?(&mut Env::new())?, Atom::atom(1)); 306 | Ok(()) 307 | } 308 | 309 | #[test] 310 | fn test_set() -> Result<(), TestError> { 311 | assert_eq!(compile_expr("(let :x 1 (do (set :x 2) (get :x)))")?(&mut Env::new())?, Atom::atom(2)); 312 | Ok(()) 313 | } 314 | 315 | #[test] 316 | fn test_factorial() -> Result<(), TestError> { 317 | assert_eq!(compile_expr(" 318 | (let :acc 1 319 | (let :i 0 320 | (loop 10 (do 321 | (set :i (+ (get :i) 1)) 322 | (if (= (get :i) 10) 323 | (get :acc) 324 | (set :acc (* (get :acc) (get :i)))) 325 | )))) 326 | ")?(&mut Env::new())?, Atom::atom(362880)); 327 | Ok(()) 328 | } 329 | 330 | #[test] 331 | fn test_nom_example() -> Result<(), TestError> { 332 | assert_eq!(compile_expr(" 333 | ((if (= (+ 3 (/ 9 3)) 334 | (* 2 3)) 335 | * 336 | /) 337 | 456 123) 338 | ")?(&mut Env::new())?, Atom::atom(56088)); 339 | Ok(()) 340 | } 341 | 342 | #[test] 343 | fn test_dynamic_function() -> Result<(), TestError> { 344 | assert_eq!(compile_expr(" 345 | (let :f '(+ 1 2) 346 | ((get :f)) 347 | ) 348 | ")?(&mut Env::new())?, Atom::atom(3)); 349 | Ok(()) 350 | } 351 | } 352 | -------------------------------------------------------------------------------- /lineiform/src/lib.rs: -------------------------------------------------------------------------------- 1 | #![allow(unused_imports, unused_parens, unused_variables)] 2 | #![deny(unused_must_use, improper_ctypes_definitions)] 3 | #![feature(box_syntax, box_patterns, trait_alias, unboxed_closures, fn_traits, ptr_metadata, stmt_expr_attributes, entry_insert, map_try_insert, if_let_guard, bench_black_box, associated_type_bounds, asm, destructuring_assignment, generic_const_exprs, inline_const, inline_const_pat, int_roundings)] 4 | 5 | #[deny(bare_trait_objects)] 6 | extern crate jemallocator; 7 | extern crate thiserror; 8 | #[macro_use] 9 | extern crate enum_display_derive; 10 | extern crate yaxpeax_arch; 11 | extern crate yaxpeax_x86; 12 | extern crate goblin; 13 | extern crate bitvec; 14 | extern crate bitflags; 15 | extern crate rangemap; 16 | extern crate tangle; 17 | extern crate frunk; 18 | use yaxpeax_x86::long_mode::{Operand, Instruction, RegSpec, Opcode}; 19 | use std::collections::HashMap; 20 | use core::marker::PhantomData; 21 | use std::sync::Arc; 22 | use std::rc::Rc; 23 | use core::num::Wrapping; 24 | use core::pin::Pin; 25 | 26 | 27 | pub mod lift; 28 | pub mod tracer; 29 | pub mod block; 30 | pub mod marker; 31 | use crate::lift::JitValue; 32 | 33 | /// Rust closures are rust-call calling convention only. This is a problem, since 34 | /// we want to lift them to our IR, and we don't know what registers are "live" 35 | /// at the start and end. We instead make it so that making a Function from an Fn 36 | /// actually starts the trace from the Trampoline call for the correct type, 37 | /// with it's call being extern "C". Rustc then emits a stdcall -> rust-call 38 | /// prologue and epilogue for us, which we can lift. 39 | 40 | extern "C" fn c_fn( 41 | cb: extern fn(data: *const c_void, A)->O, // "rust-call", but adding that ICEs 42 | d: *const c_void, 43 | a: A 44 | ) -> O { 45 | cb(d, a) 46 | } 47 | 48 | use core::ffi::c_void; 49 | #[derive(Copy, Clone)] 50 | struct FastFn { 51 | f: extern "C" fn(usize, data: *const c_void, A)->O, 52 | f_body: Pin, 53 | } 54 | 55 | impl FnOnce for FastFn { 56 | type Output=O; 57 | extern "rust-call" fn call_once(self, args: A) -> O { 58 | unimplemented!(); 59 | } 60 | } 61 | 62 | impl FnMut for FastFn where FastFn: FnOnce { 63 | extern "rust-call" fn call_mut(&mut self, args: A) -> as FnOnce>::Output { 64 | unimplemented!(); 65 | } 66 | } 67 | 68 | use core::hint::black_box; 69 | #[inline(never)] 70 | #[no_mangle] 71 | pub extern "C" fn break_here(a: usize) -> usize { 72 | black_box(a) 73 | } 74 | 75 | impl Fn for FastFn where 76 | FastFn: FnOnce, 77 | { 78 | extern "rust-call" fn call(&self, args: A) -> >::Output { 79 | println!("got arguments: {:?}", args); 80 | black_box(break_here(black_box(1))); 81 | (self.f)(0, self as *const Self as *const c_void, args) 82 | } 83 | } 84 | 85 | use lift::Jit; 86 | use tracer::Tracer; 87 | use block::Function; 88 | pub struct Lineiform { 89 | decompiled: HashMap<*const (), Arc>>, 90 | _phantom: PhantomDataO>, 91 | tracer: Tracer<'static>, 92 | } 93 | 94 | use std::mem::size_of; 95 | impl Lineiform { 96 | pub fn new() -> Self { 97 | Self { 98 | tracer: Tracer::new().unwrap(), 99 | decompiled: HashMap::new(), 100 | _phantom: PhantomData, 101 | } 102 | } 103 | 104 | pub fn speedup(&mut self, f: F) -> impl Fn(A)->O where 105 | F: Fn(A)->O, 106 | [(); usize::div_ceil(size_of::(), size_of::())]: Sized, 107 | [(); usize::div_ceil(size_of::<(extern fn(data: *const c_void, A)->O, *const c_void, A)>(), size_of::())]: Sized, 108 | [(); usize::div_ceil(size_of::() ,size_of::())]: Sized 109 | { 110 | // we pin the closure in place, because we need to be able to inline the pointer 111 | // value when jitting and thus it can never move 112 | let f_pinned = Box::pin(f); 113 | let f_size = std::mem::size_of_val(f_pinned.as_ref().get_ref()); 114 | // We are given a [data, trait] fat pointer for f. 115 | // We want to feed c_fn into our lifter, calling F::call as arg 1, with 116 | // data as our argument 2. 117 | let call: for <'a> extern "rust-call" fn(&'a F, (A,))->O = >::call; 118 | let f_body = f_pinned.as_ref().get_ref() as *const _ as *const u8; 119 | let call: extern fn(data: *const c_void, A)->O = 120 | unsafe { std::mem::transmute(call) }; 121 | // We now compile c_fn(call, f_body, a: A) to a function that can throw 122 | // away call and f_body. 123 | use crate::lift::Location; 124 | let mut func = Function::<(extern fn(data: *const c_void, A)->O, *const c_void, A), O>::new(&mut self.tracer, c_fn:: as *const ()) 125 | .unwrap() 126 | .assume_args(vec![ 127 | (Location::Reg(RegSpec::rdi()), JitValue::Const(Wrapping(call as usize))), 128 | (Location::Reg(RegSpec::rsi()), JitValue::Ref( 129 | Rc::new(JitValue::Frozen { 130 | addr: f_body, 131 | size: f_size, 132 | }), 0) 133 | ), 134 | ]); 135 | let func_size = func.size; 136 | let mut inlined = Jit::new(&mut self.tracer); 137 | //inlined.assume(vec![call as *const u8, f_body as *const u8]); 138 | let (inlined, _size) = inlined.lower::<_, O>(func).unwrap(); 139 | 140 | if true { 141 | let orig_func_dis = self.tracer.disassemble(c_fn:: as *const (), func_size as usize).unwrap(); 142 | println!("original function:"); 143 | self.tracer.format(&orig_func_dis).unwrap(); 144 | //println!("\npinned rdi={:x} rsi={:x}", call as usize, f_body as usize); 145 | //black_box(f_body); 146 | 147 | let new_func_dis = self.tracer.disassemble(inlined as *const (), _size as usize).unwrap(); 148 | println!("\nrecompiled function:"); 149 | self.tracer.format(&new_func_dis).unwrap(); 150 | } 151 | 152 | // TODO: cache this? idk what our key is 153 | FastFn { f: unsafe { std::mem::transmute(inlined) }, f_body: f_pinned } 154 | } 155 | } 156 | -------------------------------------------------------------------------------- /lineiform/src/lift.rs: -------------------------------------------------------------------------------- 1 | use crate::tracer::{self, Tracer}; 2 | use tangle::ir::PortIdx; 3 | use tangle::node::Continuation; 4 | 5 | use crate::block::{Function}; 6 | 7 | use yaxpeax_x86::long_mode::{Operand, Instruction, RegSpec, Opcode}; 8 | use yaxpeax_arch::LengthedInstruction; 9 | 10 | use bitvec::vec::BitVec; 11 | use bitvec::slice::BitSlice; 12 | 13 | use core::marker::PhantomData; 14 | use std::collections::HashMap; 15 | use std::convert::{TryInto, TryFrom}; 16 | use core::num::Wrapping; 17 | use core::cmp::max; 18 | use rangemap::RangeInclusiveMap; 19 | 20 | const HOST_WIDTH: u8 = std::mem::size_of::() as u8; 21 | const STACK_TOP: usize = 0xFFFF_FFFF_FFFF_FFF0; 22 | 23 | use thiserror::Error; 24 | #[derive(Error, Debug)] 25 | pub enum LiftError { 26 | #[error("unable to lift function")] 27 | Lifting, 28 | #[error("unsupported target")] 29 | Unsupported, 30 | #[error("unknown register {0}")] 31 | UnknownRegister(&'static str), 32 | #[error("lifting unimplemented for inst {0}")] 33 | UnimplementedInst(Instruction), 34 | #[error("disassembling call target {0}")] 35 | Block(#[from] crate::block::BlockError), 36 | #[error("tracer error {0}")] 37 | Tracer(#[from] crate::tracer::TracerError), 38 | #[error("we hit a return to {0:x}")] 39 | FunctionEnded(usize), 40 | #[error("we didn't return to the callsite in our lifted function")] 41 | TranslateOut(()), 42 | } 43 | 44 | trait OperandFun { 45 | fn is_imm(&self) -> bool; 46 | } 47 | 48 | impl OperandFun for Operand { 49 | fn is_imm(&self) -> bool { 50 | match self { 51 | Operand::ImmediateI8( _) | Operand::ImmediateI16(_) | 52 | Operand::ImmediateI32(_) | Operand::ImmediateI64(_) => 53 | true, 54 | _ => false, 55 | } 56 | } 57 | } 58 | 59 | /// We use these locations essentially as the keys of an x86 operand -> Tangle 60 | /// port map. When we create a function, we enter all the calling convention's 61 | /// inputs into a store, and when we decode x86 and lift it we try and fetch the 62 | /// "variable" that operand corrosponds to. 63 | use std::fmt::Display; 64 | #[derive(Clone, Hash, PartialEq, Eq, Debug, Display)] 65 | pub enum Location { 66 | /// A register value 67 | Reg(RegSpec), 68 | /// A stack slot value. The `usize` is the "stack address" of it: 69 | /// `push rax` will create a new Stack(STACK_TOP) entry if it's the first 70 | /// stack operation, for example. 71 | Stack(usize), 72 | } 73 | 74 | #[derive(Clone, Debug, PartialEq, Eq)] 75 | pub enum JitVariable { 76 | Variable(PortIdx), 77 | Known(JitValue), 78 | } 79 | 80 | /* 81 | impl JitVariable { 82 | /// Get a native width JitValue from a JitVariable. You should probably not 83 | /// use this. 84 | pub fn value(&mut self, builder: &mut FunctionBuilder) -> JitValue { 85 | match self { 86 | JitVariable::Variable(var) => { 87 | let mut val = builder.use_var(*var); 88 | JitValue::Value(val) 89 | }, 90 | JitVariable::Known(val) => { 91 | val.clone() 92 | }, 93 | } 94 | } 95 | 96 | /// Get a proper width temporary JitTemp value from this JitVariable. 97 | /// In order to store it in the concrete context as a JitValue, it must be 98 | /// either zero or sign extended. 99 | pub fn tmp(&mut self, builder: &mut FunctionBuilder, width: u8) -> JitTemp { 100 | match self { 101 | JitVariable::Variable(var) => { 102 | let mut val = builder.use_var(*var); 103 | use types::{I8, I16, I32, I64}; 104 | let typ = match width { 105 | 1 => I8, 106 | 2 => I16, 107 | 4 => I32, 108 | 8 => I64, 109 | _ => unimplemented!(), 110 | }; 111 | if width != HOST_WIDTH { 112 | // we want a memory or stack value as a certain width: 113 | // get a use of the variable backing it, and mask to the correct 114 | // size. 115 | JitTemp::Value(builder.ins().ireduce(typ, val), typ) 116 | } else { 117 | JitTemp::Value(val, typ) 118 | } 119 | }, 120 | JitVariable::Known(val) => { 121 | //let width_mask = Wrapping((1 << ((width*8)+1)) - 1); 122 | match val { 123 | JitValue::Const(c) => match width { 124 | 1 => JitTemp::Const8(c.0 as u8), 125 | 2 => JitTemp::Const16(c.0 as u16), 126 | 4 => JitTemp::Const32(c.0 as u32), 127 | 8 => JitTemp::Const64(c.0 as u64), 128 | _ => unimplemented!(), 129 | }, 130 | id if width == HOST_WIDTH => match id { 131 | JitValue::Ref(base, offset) => 132 | JitTemp::Ref(base.clone(), *offset), 133 | JitValue::Frozen { addr, size } => 134 | JitTemp::Frozen { addr: *addr, size: *size }, 135 | x @ _ => unimplemented!("temp {:?}", x), 136 | }, 137 | _ => unimplemented!(), 138 | } 139 | }, 140 | } 141 | } 142 | } 143 | */ 144 | 145 | use std::rc::Rc; 146 | #[derive(Debug, Clone, PartialEq, Eq)] 147 | pub enum JitValue { 148 | /// An SSA value 149 | Value(PortIdx), 150 | /// An SSA comparison flag value 151 | Flag(PortIdx), 152 | /// A reference to a value, plus an offset into it 153 | Ref(Rc, usize), 154 | /// A memory region we use to represent the stack: `sub rsp, 0x8` becomes 155 | /// manipulating a Ref to this, and load/stores can be concretized. 156 | Stack, 157 | /// A frozen memory region. We can inline values from it safely. 158 | Frozen { addr: *const u8, size: usize }, 159 | /// A statically known value. 160 | Const(Wrapping), 161 | } 162 | 163 | impl JitValue { 164 | } 165 | 166 | 167 | // TODO: fp registers and stuff 168 | type Type = usize; 169 | 170 | /// These are the temporary values used by operations - they, unlike JitValues, 171 | /// are variable width, and have operations to sign-extend or zero-extend that 172 | /// yield native width JitValues. 173 | /// This also allows us to optimize a bit, since we can leave off a few masking 174 | /// operations for values if we know exactly what width they are. 175 | #[derive(Debug, Clone, PartialEq, Eq)] 176 | pub enum JitTemp { 177 | /// An SSA value 178 | Value(PortIdx, Type), 179 | Flag(PortIdx), 180 | // XX: how do we handle these "opaque" values? we want to be able to reassemble 181 | // them from a memcpy, for example. worry about it when we have a better 182 | // use-case? 183 | /// A reference to a value, plus an offset into it 184 | Ref(Rc,usize), 185 | /// A memory region we use to represent the stack: `sub rsp, 0x8` becomes 186 | /// manipulating a Ref to this, and load/stores can be concretized. 187 | //Stack, 188 | /// A frozen memory region. We can inline values from it safely. 189 | Frozen { addr: *const u8, size: usize }, 190 | /// A statically known value. 191 | Const8(u8), 192 | Const16(u16), 193 | Const32(u32), 194 | Const64(u64), 195 | } 196 | 197 | #[derive(Clone)] 198 | struct Context { 199 | /// Our currently known execution environment. This holds values like `rax` 200 | /// and stored stack values, mapping them to either constant propagated 201 | /// values, or Cranelift SSA variables if they are unknown at compilation 202 | /// time. 203 | bound: HashMap, 204 | flags: Vec, // array of eflags source ports 205 | stack: usize, 206 | } 207 | 208 | impl Context { 209 | fn new() -> Self { 210 | Self { 211 | bound: [(Location::Reg(RegSpec::rsp()), JitVariable::Known(JitValue::Const(Wrapping(STACK_TOP))))].into(), 212 | flags: Vec::new(), 213 | stack: STACK_TOP 214 | } 215 | } 216 | 217 | fn stack(&self) -> usize { 218 | let sp = self.bound.get(&Location::Reg(RegSpec::rsp())); 219 | match sp { 220 | Some(JitVariable::Known(JitValue::Ref(base, offset))) 221 | if let JitValue::Stack = **base => 222 | { 223 | //println!("current stack @ 0x{:x}", offset); 224 | assert!(offset <= &STACK_TOP); 225 | *offset 226 | }, 227 | Some(JitVariable::Known(JitValue::Const(Wrapping(stack)))) => { 228 | assert!(stack <= &STACK_TOP); 229 | *stack 230 | } 231 | x => unimplemented!("non-concrete stack?? {:?}", x), 232 | } 233 | } 234 | 235 | fn dump(&self) { 236 | } 237 | } 238 | 239 | pub struct Jit<'a> { 240 | tracer: &'a mut Tracer<'static>, 241 | } 242 | 243 | #[derive(Debug)] 244 | enum EmitEffect { 245 | Advance, 246 | Jmp(JitValue), 247 | Branch(PortIdx, PortIdx, usize), 248 | Call(usize), 249 | Ret(Option), 250 | } 251 | 252 | struct JitFunction { 253 | ir: tangle::IR, 254 | f: tangle::node::Function, 255 | context: Context, 256 | } 257 | 258 | use std::mem::size_of; 259 | impl<'a> Jit<'a> { 260 | pub fn new(tracer: &'a mut Tracer<'static>) -> Jit<'a> { 261 | 262 | Self { 263 | tracer, 264 | } 265 | } 266 | 267 | pub fn lower(&mut self, f: Function) 268 | -> Result<(extern "C" fn((A,))->O, usize),LiftError> where 269 | [(); usize::div_ceil(size_of::(), size_of::())]: Sized, 270 | [(); usize::div_ceil(size_of::(), size_of::())]: Sized { 271 | //assert_eq!(f.pinned, vec![]); 272 | let mut ir = tangle::IR::new(); 273 | let ctx = Context::new(); 274 | println!("lowering function {}->{}", size_of::(), size_of::()); 275 | let mut tangle_f = tangle::node::Node::function::<{ usize::div_ceil(size_of::(), size_of::()) }, { usize::div_ceil(size_of::() ,size_of::()) }>(&mut ir); 276 | let mut func = JitFunction { 277 | ir: ir, 278 | f: tangle_f, 279 | context: ctx, 280 | }; 281 | let translate_result = func.translate(f, self)?; 282 | 283 | func.ir.set_body(func.f); 284 | Ok(func.ir.compile_fn().unwrap()) 285 | } 286 | } 287 | 288 | impl JitFunction { 289 | /// Translate a Function into Cranelift IR. 290 | fn translate<'a, A, O> 291 | (&mut self, f: Function, jit: &mut Jit<'a>) -> Result<(), LiftError> { 292 | println!("translate for {}->{}", A_n, O_n); 293 | // add and bind ports for all the function arguments 294 | use tangle::abi::{AbiRequest, AbiStorage}; 295 | let (abi_arg, abi_ret) = self.ir.in_region(self.f.region, |r, ir| { 296 | let abi = r.abi.as_ref().unwrap(); 297 | ( 298 | abi.provide_arguments(vec![AbiRequest::Integer(A_n)]), 299 | abi.provide_returns(vec![AbiRequest::Integer(O_n)]) 300 | ) 301 | }); 302 | for i in 0..A_n { 303 | let p = self.f.add_argument(&mut self.ir); 304 | println!("adding argument {} -> {:?}", i, p); 305 | match abi_arg[i] { 306 | AbiStorage::Register(r) => self.store(Operand::Register(r), p, None), 307 | _ => unimplemented!(), 308 | } 309 | // TODO 310 | } 311 | 312 | // populate pinned values 313 | for (place, pin) in f.pinned { 314 | println!("pinning {}", place); 315 | let oper_place = match place { 316 | Location::Reg(r) => Operand::Register(r), 317 | _ => unimplemented!() 318 | }; 319 | match pin { 320 | JitValue::Const(c) => { 321 | self.context.bound.insert(place, JitVariable::Known(JitValue::Const(c))); 322 | //let c = self.constant(c.0 as isize); 323 | //self.store(oper_place, c, None); 324 | }, 325 | JitValue::Ref(r, o) if let JitValue::Frozen { addr, size} = *r => { 326 | // for now just set it to a constant for the pointer 327 | self.context.bound.insert(place, JitVariable::Known(JitValue::Const( 328 | Wrapping((addr as isize + o as isize) as usize) 329 | ))); 330 | //let c = self.constant(addr as isize + o as isize); 331 | //self.store(oper_place, c, None); 332 | }, 333 | x => unimplemented!("pinned value {:?} = {:?}", oper_place, x) 334 | } 335 | } 336 | 337 | // jit all the basic blocks 338 | let mut need = vec![JitPath(f, f.region, variables)]; 339 | while let Some(next) = need.pop() { 340 | let mut new = function_translator.translate(next)?; 341 | need.append(&mut new); 342 | } 343 | 344 | // add function return values 345 | if self.f.cont == Continuation::Return { 346 | for i in 0..O_n { 347 | let p = self.f.add_return(&mut self.ir); 348 | println!("adding return {} -> {:?}", i, p); 349 | match abi_ret[i] { 350 | AbiStorage::Register(r) => { 351 | let reg = self.port_for_register(ip, r); 352 | self.ir.in_region(self.f.region,|r, ir| { 353 | r.connect_ports(reg, p) 354 | }); 355 | }, 356 | _ => unimplemented!(), 357 | } 358 | // TODO 359 | } 360 | } 361 | 362 | Ok(()) 363 | } 364 | } 365 | 366 | struct FunctionTranslator<'a, A, O> { 367 | ir: &'a mut IR, 368 | context: Context, 369 | region: RegionIdx, 370 | f: &'a mut Function, 371 | } 372 | 373 | pub struct JitPath(Function, RegionIdx, Context); 374 | 375 | impl FunctionTranslator<'a, A, O> { 376 | fn translate(&mut self, JitPath(f, blk, ctx): JitPath) -> Result<(), LiftError> { 377 | self.context = ctx; 378 | self.region = blk; 379 | let mut base = (f.base as usize) + f.offset.0; 380 | let mut ip = base; 381 | println!("starting emitting @ base+0x{:x}", ip - jit.tracer.get_base()?); 382 | let mut stream = f.instructions.clone(); 383 | let mut start_i = f.offset.1; 384 | let mut i = f.offset.1; 385 | 386 | loop { 387 | // snapshot of our context when we started this basic block 388 | let mut snap = self.context.bound.clone(); 389 | let inst = stream.get(i); 390 | if let None = inst { 391 | // we ran out of inputs but didn't hit a ret! 392 | println!("hit translate out @ 0x{:x}", ip); 393 | return Err(LiftError::TranslateOut(())); 394 | } 395 | let inst = inst.unwrap(); 396 | ip += inst.len().to_const() as usize; 397 | let mut eff = self.emit(ip, inst)?; 398 | match eff { 399 | EmitEffect::Advance => { 400 | i += 1; 401 | }, 402 | EmitEffect::Jmp(p) => { 403 | if let JitValue::Const(c) = p { 404 | // TODO: for now we only care about tailcalls; if it's 405 | // inside the function region then it's internal flow control 406 | // and we will handle it later once we implement RVSDG 407 | // irreducable cfg; we will analyze the cfg before lifting 408 | // and build branches+tail controlled loop nodes for internal 409 | // jumps. 410 | // a jump table with an internal jump label might not work 411 | // with that? it would probably just turn into something 412 | // duplicated and not break, though. 413 | if target >= (f.base as usize) && target < (f.base as usize) + f.size { 414 | unimplemented!("internal jmp"); 415 | } 416 | return 417 | } 418 | let mut tailcall = tangle::node::Node::leave(); 419 | 420 | // create input ports for all values that need to be alive for the 421 | // jmp 422 | use tangle::port::Storage; 423 | let reg_ports = snap.iter() 424 | .filter(|(bound, _)| **bound != Location::Reg(RegSpec::rsp())) 425 | .map(|(bound, alive)| { 426 | println!("materializing {} <- {:?} for jmp", bound, alive); 427 | match alive { 428 | JitVariable::Variable(p) => { 429 | self.ir.in_region(self.f.region, |r, ir| { 430 | let material = r.add_port(); 431 | match bound { 432 | Location::Reg(bound_reg) => 433 | r.constrain(material, Storage::Physical(*bound_reg)), 434 | Location::Stack(ss) => 435 | // stack slots just need to observe the 436 | // state edge 437 | r.constrain(material, Storage::Immaterial(None)), 438 | } 439 | (*p, material) 440 | }) 441 | }, 442 | JitVariable::Known(JitValue::Const(c)) => { 443 | let c = self.constant(c.0 as isize); 444 | 445 | self.ir.in_region(self.f.region, |r, ir| { 446 | let material = r.add_port(); 447 | match bound { 448 | Location::Reg(bound_reg) => 449 | r.constrain(material, Storage::Physical(*bound_reg)), 450 | Location::Stack(ss) => 451 | // stack slots just need to observe the 452 | // state edge 453 | r.constrain(material, Storage::Immaterial(None)), 454 | } 455 | (c, material) 456 | }) 457 | }, 458 | x => unimplemented!("jmp materialize {:?}", x), 459 | } 460 | }).collect::>(); 461 | 462 | // add the tailcall node to the graph, and also bind as inputs 463 | // all the alive ports 464 | let (_, tailcall_o) = self.ir.in_region(self.region, &mut self.ir, |tailcall, ir| { 465 | tailcall.connect_input(0, p, ir); 466 | for (i, (reg_port, input_port)) in reg_ports.iter().enumerate() { 467 | tailcall.add_input(*input_port, ir); 468 | tailcall.connect_input(1+i, *reg_port, ir); 469 | } 470 | tailcall.outputs[0] 471 | }); 472 | self.f.cont = Continuation::Jmp(tailcall_o); 473 | break; 474 | }, 475 | EmitEffect::Ret(targ) => { 476 | if let Some(targ) = targ { 477 | unimplemented!("return to known"); 478 | } else { 479 | assert_eq!(self.f.cont, Continuation::Return); 480 | break; 481 | } 482 | }, 483 | x => unimplemented!("emit effect {:?}", x), 484 | } 485 | } 486 | 487 | } 488 | 489 | fn emit(&mut self, ip: usize, inst: &Instruction) -> Result { 490 | println!("---- {}", inst); 491 | let ms = inst.mem_size().and_then(|m| m.bytes_size()); 492 | match inst.opcode() { 493 | Opcode::MOV => { 494 | // If we have a mov that purely swaps registers, we can simply 495 | // copy over the values instead of having to introduce an extraneous 496 | // mov; this may save us work when doing optimizations. 497 | match (inst.operand(0), inst.operand(1)) { 498 | (Operand::Register(r0), Operand::Register(r1)) => { 499 | println!("simple mov"); 500 | self.context.bound.insert(Location::Reg(r0), 501 | self.context.bound.get(&Location::Reg(r1)).unwrap().clone()); 502 | return Ok(EmitEffect::Advance) 503 | }, 504 | _ => () 505 | }; 506 | // we need to support mov dword [eax+4], rcx and vice versa, so 507 | // both the load and store need a mem_size for memory width. 508 | let val = self.port_for(ip, inst.operand(1), ms); 509 | self.store(inst.operand(0), val, ms); 510 | }, 511 | //Opcode::ADD => { 512 | // let left = self.operand_value(ip, inst.operand(0), ms); 513 | // let mut right = self.operand_value(ip, inst.operand(1), ms); 514 | // if inst.operand(1).is_imm() { 515 | // right = right.sign_extend(self.builder, inst.operand(0).width().or(ms).unwrap()); 516 | // } 517 | // let added = self.add::(left, right); 518 | // self.store(inst.operand(0), added, ms); 519 | //}, 520 | //Opcode::ADC => { 521 | // let left = self.operand_value(ip, inst.operand(0), ms); 522 | // let mut right = self.operand_value(ip, inst.operand(1), ms); 523 | // if inst.operand(1).is_imm() { 524 | // right = right.sign_extend(self.builder, inst.operand(0).width().or(ms).unwrap()); 525 | // } 526 | // let carry = self.context.check_flag(Flags::CF, true, self.builder); 527 | // println!("adc {:?} {:?} {:?}", left, right, carry); 528 | // let added = self.adc::(left, right, carry); 529 | // self.store(inst.operand(0), added, ms); 530 | //}, 531 | op @ (Opcode::SUB | Opcode::CMP) => { 532 | unimplemented!(); 533 | if let (Opcode::SUB, Operand::Register(const { RegSpec::rsp() })) = (op, inst.operand(0)) { 534 | // This is a sub rsp, 8; or whatever instruction allocating stack 535 | // space. We don't want to emit any instructions, just allocate 536 | // stack slots. 537 | let space = match inst.operand(1) { 538 | Operand::ImmediateI8(i) => usize::try_from(i).unwrap(), 539 | x => unimplemented!("unknown space for stack allocating {:?}", x), 540 | }; 541 | println!("emitting {} stack slots", space / HOST_WIDTH as usize); 542 | assert!(space % HOST_WIDTH as usize == 0); 543 | assert!(self.context.stack() == STACK_TOP); 544 | unimplemented!("shift_stack for sub"); 545 | for i in 0..(space / HOST_WIDTH as usize) { 546 | self.f.add_stack_slot(&mut self.ir); 547 | } 548 | } else { 549 | // sub eax, ff 550 | // ff gets sign extended to 0xffff_ffff 551 | let left = self.port_for(ip, inst.operand(0), ms); 552 | let mut right = self.port_for(ip, inst.operand(1), ms); 553 | unimplemented!() 554 | //if inst.operand(1).is_imm() { 555 | // right = right.sign_extend(self.builder, inst.operand(0).width().or(ms).unwrap()); 556 | //} 557 | //let subbed = self.sub::(left, right); 558 | //if let Opcode::SUB = op { 559 | // self.store(inst.operand(0), subbed, ms); 560 | //} 561 | } 562 | }, 563 | Opcode::JMP => { 564 | // We use jump_target, which calls operand_value, because if we 565 | // can get a JitValue::Const for jump target immediately we should. 566 | // This allows us to immediately start JITting tailcall closures 567 | // without having to stop lifting, do an optimization pass, and 568 | // then notice we have a jmp with a constant operand and lift 569 | // again in the optimizer; that leads to a non-trivial amount of 570 | // overhead, including emitting a lot of unneeded ports and edges 571 | // for a Leave instruction. 572 | let target = self.jump_target(inst.operand(0), ip, ms); 573 | // If we know where the jump is going, we can try to inline 574 | if let JitValue::Const(c) = target { 575 | println!("known jump location: 0x{:x}", c.0); 576 | return Ok(EmitEffect::Jmp(target)); 577 | } else if let JitValue::Value(v) = target { 578 | return Ok(EmitEffect::Jmp(target)); 579 | } 580 | }, 581 | Opcode::LEA => { 582 | let val = self.effective_address(ip, inst.operand(1))?; 583 | self.store(inst.operand(0), val, inst.operand(1).width()); 584 | }, 585 | Opcode::RETURN => { 586 | if inst.operand_present(0) { 587 | unimplemented!("return with stack pops"); 588 | } 589 | // if we were trying to return to the value past the top of the 590 | // stack, then we were returning out of our emitted function: 591 | // we're all done jitting, and can tell cranelift to stop. 592 | if self.context.stack() == STACK_TOP { 593 | return Ok(EmitEffect::Ret(None)); 594 | } 595 | let new_rip = self.context.bound.get(&Location::Reg(RegSpec::rsp())); 596 | if let Some(&JitVariable::Known(JitValue::Const(c))) = new_rip { 597 | self.shift_stack(1); 598 | return Ok(EmitEffect::Ret(Some(c.0))); 599 | } else { 600 | self.context.dump(); 601 | unimplemented!("return to unknown location {:?}", new_rip); 602 | } 603 | }, 604 | x => unimplemented!("unimplemented opcode {:?}", x), 605 | } 606 | Ok(EmitEffect::Advance) 607 | } 608 | 609 | /// Get the Tangle port that corrosponds with the requested operand. 610 | /// 611 | /// If we have `mov eax, ecx`, we want to resolve ecx to the last emitted ecx port. 612 | /// This has to handle *all* operand types, however: if there's an `[ecx]` we need 613 | /// to emit a StackLoad instruction and return the value result of that as well. 614 | /// 615 | /// If the location is a load, `width` is the memory operation size in bytes. 616 | /// It is ignored for other operand types, so that e.g. `mov dword [rax+4], rcx` 617 | /// doesn't mistakenly truncate rcx. 618 | pub fn port_for(&mut self, ip: usize, operand: Operand, ms: Option) -> PortIdx { 619 | println!("getting port for {}", operand); 620 | match operand { 621 | Operand::RegDisp(const { RegSpec::rsp() }, o) => { 622 | unimplemented!("stack ref"); 623 | }, 624 | Operand::Register(const { RegSpec::rip() }) => { 625 | unimplemented!("emit const") 626 | //HOST_TMP(ip as _) 627 | }, 628 | //Operand::RegDisp(const { RegSpec::rip() }, o) => { 629 | // HOST_TMP((Wrapping(ip) + Wrapping(o as isize as usize)).0 as _) 630 | //}, 631 | //Operand::Register(r) if r.width() == 1 => { 632 | // match r { 633 | // const { RegSpec::al() } | const { RegSpec::bl() } | 634 | // const { RegSpec::cl() } => { 635 | // // I should probably re-think Const, and add variable 636 | // // widths to it: there's not really a good reason for 637 | // // `and al, 0x1` to have to mask out rax for the value 638 | // // and then bitselect it back in at the store. 639 | // self.get(Location::Reg(r)).tmp(self.builder, r.width()) 640 | // }, 641 | // _ => unimplemented!("al etc unimplemented") 642 | // } 643 | //}, 644 | Operand::Register(r) => { 645 | self.port_for_register(ip, r) 646 | }, 647 | x => unimplemented!("getting a port for {} unimplemented", x), 648 | } 649 | } 650 | 651 | /// Store the value corrosponding with the given PortIdx in the storage location of the 652 | /// operand, using the required memory width. 653 | pub fn store(&mut self, operand: Operand, val: PortIdx, ms: Option) { 654 | println!("storing port in {} width {:?}", operand, ms); 655 | match operand { 656 | Operand::Register(const { RegSpec::rip() }) => { 657 | unimplemented!("store to ip to jump"); 658 | }, 659 | Operand::Register(const { RegSpec::rsp() }) => { 660 | unimplemented!("store to stack register"); 661 | }, 662 | Operand::Register(r) if r.width() == HOST_WIDTH => { 663 | self.context.bound.insert(Location::Reg(r), JitVariable::Variable(val)) 664 | }, 665 | _ => unimplemented!("store to unimplemented location {:?}", operand) 666 | }; 667 | } 668 | 669 | fn jump_target(&mut self, op: Operand, ip: usize, ms: Option) -> JitValue { 670 | // Resolve the jump target (either absolute or relative) 671 | let target: JitValue = match op { 672 | absolute @ ( 673 | Operand::Register(_) | 674 | Operand::RegDeref(_) | 675 | Operand::RegDisp(_, _) 676 | ) => { 677 | self.operand_value(absolute, ip, ms) 678 | }, 679 | relative @ ( 680 | Operand::ImmediateI8(_) | 681 | Operand::ImmediateI16(_) | 682 | Operand::ImmediateI32(_) 683 | )=> { 684 | let val = match relative { 685 | Operand::ImmediateI8(i) => i as isize as usize, 686 | Operand::ImmediateI16(i) => i as isize as usize, 687 | Operand::ImmediateI32(i) => i as isize as usize, 688 | _ => unreachable!(), 689 | }; 690 | JitValue::Const(Wrapping(ip) + Wrapping(val)) 691 | }, 692 | _ => unimplemented!("jump target operand {}", op), 693 | }; 694 | target 695 | } 696 | 697 | /// Get the JitValue of an operand, of the correct width; if it can be represented 698 | /// as a JitValue::Const attempt to do so. This function should be used sparingly 699 | /// since baking too many constants in a JIT method may be inefficient. 700 | fn operand_value(&mut self, op: Operand, ip: usize, ms: Option) -> JitValue { 701 | let src = match op { 702 | Operand::Register(r) => self.context.bound.get(&Location::Reg(r)).unwrap(), 703 | _ => todo!(), 704 | }; 705 | match src { 706 | JitVariable::Known(val) => val.clone(), 707 | x => unimplemented!("{:?}", x) 708 | } 709 | } 710 | 711 | /// Return a port representing the current live value of the required register 712 | fn port_for_register(&mut self, ip: usize, reg: RegSpec) -> PortIdx { 713 | assert!(reg != RegSpec::rip()); 714 | assert!(reg != RegSpec::rsp()); 715 | match self.context.bound.entry(Location::Reg(reg)).or_insert_with(|| { 716 | // We only get the variable on reads, and if we are reading from 717 | // an unbound register than we have to create a port and bind it to 718 | // the register as physical storage. 719 | // This is because hypothetically we could get a Rust ABI function 720 | // that reads one of the registers that isn't actually from an argument 721 | // (though it would always be junk, so do we care to be correct about it?) 722 | panic!("unset register {:?}", reg); 723 | JitVariable::Variable(self.f.add_argument(&mut self.ir)) 724 | }) { 725 | JitVariable::Variable(v) => *v, 726 | JitVariable::Known(c) => unimplemented!("constant pool"), 727 | } 728 | } 729 | 730 | fn effective_address(&mut self, ip: usize, loc: Operand) -> Result { 731 | assert_eq!(HOST_WIDTH, 8); 732 | match loc { 733 | Operand::RegDeref(r) => { 734 | unimplemented!(); 735 | }, 736 | Operand::RegDisp(r, disp) => { 737 | // we can't do this, because `lea bl, [ah]` would be wrong. 738 | // self::get is kinda dangerous huh 739 | //let reg = self.get(Location::Reg(r)).tmp(self.builder, r.width()); 740 | let reg = self.port_for_register(ip, r); 741 | //let zx = reg.zero_extend(self.builder, HOST_WIDTH); 742 | //Ok(self.add::(zx, 743 | // HOST_TMP(disp as isize as usize as _))) // XXX: is this right? 744 | let disp_const = self.constant(disp.try_into().unwrap()); 745 | Ok(self.add(reg, disp_const)) 746 | } 747 | _ => unimplemented!() 748 | } 749 | } 750 | 751 | /// Move the stack either up or down N slots. -1 becomes a -4 or -8 752 | pub fn shift_stack(&mut self, slots: isize) { 753 | let stack = self.context.stack(); 754 | if slots < 0 { 755 | // we are potentially creating stack slots 756 | for slot in 0..(slots*-1) { 757 | let slot = (STACK_TOP - (slot as usize*HOST_WIDTH as usize)); 758 | self.context.bound.entry(Location::Stack(slot)).or_insert_with(|| { 759 | println!("creating stack slot {:x}", slot); 760 | // bind the state variable port 761 | JitVariable::Variable(self.f.add_stack_slot(&mut self.ir)) 762 | }); 763 | } 764 | } 765 | self.context.bound.insert(Location::Reg(RegSpec::rsp()), 766 | JitVariable::Known(JitValue::Const(Wrapping( 767 | (self.context.stack() as isize + (slots*HOST_WIDTH as isize)) as usize)))); 768 | } 769 | 770 | fn constant(&mut self, c: isize) -> PortIdx { 771 | // TODO: constant pooling? idk 772 | println!("adding constant node 0x{:x}", c); 773 | let mut c = tangle::node::Node::constant(c); 774 | self.f.add_body(c, &mut self.ir, |c, r| { c.outputs[0] }).1 775 | } 776 | 777 | // TODO: correct width of operands and output 778 | fn add(&mut self, left: PortIdx, right: PortIdx) -> PortIdx { 779 | let mut add_node = tangle::node::Node::simple(tangle::node::Operation::Add); 780 | self.f.add_body(add_node, &mut self.ir, |add_node, r| { 781 | add_node.connect_input(0, left, r); 782 | add_node.connect_input(1, right, r); 783 | add_node.outputs[0] 784 | }).1 785 | } 786 | 787 | pub fn format(&self) { 788 | // FIXME: do some sort of debug printing 789 | } 790 | } 791 | -------------------------------------------------------------------------------- /lineiform/src/marker.rs: -------------------------------------------------------------------------------- 1 | trait Mark {} 2 | 3 | impl Mark for usize {} 4 | 5 | pub fn something() { 6 | } 7 | -------------------------------------------------------------------------------- /lineiform/src/tracer.rs: -------------------------------------------------------------------------------- 1 | // yaxpeax decoder example 2 | mod decoder { 3 | use yaxpeax_arch::{Arch, AddressDisplay, Decoder, Reader, ReaderBuilder}; 4 | 5 | pub fn decode_stream< 6 | 'data, 7 | A: yaxpeax_arch::Arch, 8 | U: ReaderBuilder, 9 | >(data: U) -> Vec 10 | where A::Instruction: std::fmt::Display 11 | { 12 | let mut reader = ReaderBuilder::read_from(data); 13 | let mut address: A::Address = reader.total_offset(); 14 | 15 | let decoder = A::Decoder::default(); 16 | let mut decode_res = decoder.decode(&mut reader); 17 | let mut res = Vec::new(); 18 | loop { 19 | match decode_res { 20 | Ok(inst) => { 21 | //println!("{}: {}", address.show(), inst); 22 | decode_res = decoder.decode(&mut reader); 23 | address = reader.total_offset(); 24 | res.push(inst); 25 | } 26 | Err(e) => { 27 | //println!("{}: decode error: {}", address.show(), e); 28 | break; 29 | } 30 | } 31 | } 32 | res 33 | } 34 | } 35 | 36 | use yaxpeax_x86::amd64::{Arch as x86_64}; 37 | use yaxpeax_arch::{ReaderBuilder, U8Reader}; 38 | use yaxpeax_x86::long_mode::Instruction; 39 | 40 | // Getting the size of a closure in Rust is normally impossible. 41 | // We do it by using some section hacks: 42 | // 1) We have a variable ("NEEDLE") in a base section using #[link] 43 | // 2) We open /proc/self/exe to get our currently running executable 44 | // 3) ...and then fetch the sections for it, and use that to get the virtual addr 45 | // of NEEDLE to calculate the base_addr 46 | // 4) Using the base_addr we can search our symbols for a virtual address that 47 | // matches the symbol for the closure we are looking for 48 | // 5) ...and using that symbol, we can get the size of the closure, for proper 49 | // decompilation. 50 | // Note that this depends on us having symbols 51 | 52 | use thiserror::Error; 53 | #[derive(Error, Debug)] 54 | pub enum TracerError { 55 | #[error("goblin parsing error")] 56 | Goblin(#[from] goblin::error::Error), 57 | #[error("io error")] 58 | Io(#[from] std::io::Error), 59 | #[error("unsupported platform: bluetech only supports linux")] 60 | Unsupported, 61 | #[error("can't find symbol for function {0:x}")] 62 | CantFindFunction(usize, usize), 63 | #[error("executable has no .needle section")] 64 | NoNeedle, 65 | #[error("couldn't write formatting")] 66 | Format(#[from] std::fmt::Error), 67 | } 68 | 69 | use std::marker::PhantomPinned; 70 | use std::collections::HashMap; 71 | use std::sync::Arc; 72 | pub struct Tracer<'a> { 73 | buf: &'a [u8], 74 | elf: Elf<'a>, 75 | base: Option, 76 | /// Map for cacheing decompilation of functions 77 | cache: HashMap<*const (), Arc>>, 78 | _pin: std::marker::PhantomPinned, 79 | } 80 | 81 | #[link_section = ".needle"] 82 | pub static NEEDLE: usize = 0xBADC0DE; 83 | 84 | use goblin::elf::Elf; 85 | use goblin::elf::sym::Sym; 86 | use goblin::Object; 87 | use core::pin::Pin; 88 | use std::convert::TryInto; 89 | impl<'a> Tracer<'a> { 90 | // this should be a lazy_static 91 | pub fn new() -> Result, TracerError> { 92 | use std::fs; 93 | use std::path::Path; 94 | let path = Path::new("/proc/self/exe"); 95 | let buffer = Box::leak(box fs::read(path)?); 96 | let mut t = Tracer { 97 | buf: buffer, 98 | elf: get_elf(buffer)?, 99 | base: None, 100 | cache: HashMap::new(), 101 | _pin: PhantomPinned, 102 | }; 103 | Ok(t) 104 | } 105 | 106 | // It's either this or make everyone use a linker script, which sucks 107 | pub fn get_base(&mut self) -> Result { 108 | if let Some(b) = self.base { 109 | return Ok(b); 110 | } 111 | let needle = &NEEDLE as *const usize as usize; 112 | let needle_offset = self.elf.section_headers.iter() 113 | .filter(|section| { 114 | let name = self.elf.shdr_strtab.get_at(section.sh_name); 115 | name == Some(".needle") 116 | }).next().ok_or(TracerError::NoNeedle)?.sh_addr as usize; 117 | let base = needle - needle_offset; 118 | self.base = Some(base); 119 | Ok(base) 120 | } 121 | 122 | /// Get symbol for vaddr. The vaddr may be in the middle of a symbol. 123 | pub fn get_symbol_from_vaddr(&self, f: usize) -> Option { 124 | let addr = (f as usize).try_into().unwrap(); 125 | self.elf.syms.iter().filter(|sym| 126 | (sym.st_value <= addr) && (addr < (sym.st_value + sym.st_size)) 127 | ).next() 128 | } 129 | 130 | /// Get symbol containing an address. 131 | pub fn get_symbol_from_address(&mut self, f: *const ()) -> Result { 132 | //let mut fi = std::fs::File::open("/proc/self/maps").unwrap(); 133 | //let out = std::io::stdout(); 134 | //std::io::copy(&mut fi, &mut out.lock()).unwrap(); 135 | println!("looking for {:x}", f as usize); 136 | 137 | let base = self.get_base()?; 138 | let sym = self.get_symbol_from_vaddr(f as usize - base) 139 | .ok_or(TracerError::CantFindFunction(base, f as usize - base))?; 140 | Ok(sym) 141 | } 142 | 143 | pub fn disassemble(&mut self, f: *const (), f_size: usize) -> Result>, TracerError> { 144 | let cache_hit = self.cache.get(&f); 145 | if let Some(cache_hit) = cache_hit { 146 | // Our function was in the cache - return our cached decompilation 147 | return Ok(cache_hit.clone()); 148 | } 149 | let all = decoder::decode_stream::(unsafe { 150 | core::slice::from_raw_parts( 151 | f as *const u8, 152 | f_size as usize) as &[u8] 153 | }); 154 | let entry = Arc::new(all); 155 | self.cache.insert(f, entry.clone()); 156 | Ok(entry) 157 | } 158 | 159 | pub fn format(&self, instructions: &Vec) -> Result<(), TracerError> { 160 | let mut fmt = String::new(); 161 | for instruction in instructions { 162 | instruction.write_to(&mut fmt)?; 163 | println!("{}", fmt); 164 | fmt.clear(); 165 | } 166 | Ok(()) 167 | } 168 | } 169 | 170 | fn get_elf<'a>(buf: &'a [u8]) -> Result, TracerError> { 171 | if let Object::Elf(elf) = Object::parse(buf)? { 172 | Ok(elf) 173 | } else { 174 | Err(TracerError::Unsupported) 175 | } 176 | } 177 | use core::hint::black_box; 178 | 179 | use core::num::Wrapping; 180 | #[inline(never)] 181 | pub fn mul_two(u: Wrapping) -> Wrapping { 182 | u * Wrapping(2) 183 | } 184 | #[inline(never)] 185 | pub fn add_one(u: Wrapping) -> Wrapping { 186 | mul_two(u + Wrapping(1)) + Wrapping(2) 187 | } 188 | 189 | #[cfg(test)] 190 | mod test { 191 | use crate::tracer::Tracer; 192 | use crate::tracer::TracerError; 193 | #[test] 194 | fn can_find_base() -> Result<(), TracerError> { 195 | assert_eq!(Tracer::new()?.get_base()? != 0, true); 196 | Ok(()) 197 | } 198 | 199 | #[test] 200 | fn can_resolve_function() -> Result<(), TracerError> { 201 | assert_eq!(Tracer::new()?.get_symbol_from_address(crate::tracer::add_one as *const ()) 202 | .is_ok(), true); 203 | Ok(()) 204 | } 205 | 206 | #[test] 207 | fn can_disassemble_fn() -> Result<(), TracerError> { 208 | let mut tracer = Tracer::new()?; 209 | let sym = tracer.get_symbol_from_address(crate::tracer::add_one as *const ())?; 210 | let instructions = tracer.disassemble(crate::tracer::add_one as *const (), sym.st_size as usize)?; 211 | tracer.format(&instructions)?; 212 | Ok(()) 213 | } 214 | } 215 | -------------------------------------------------------------------------------- /src/main.rs: -------------------------------------------------------------------------------- 1 | #![allow(unused_imports, unused_parens)] 2 | #![deny(unused_must_use, improper_ctypes_definitions)] 3 | #![feature(box_syntax, box_patterns, trait_alias, unboxed_closures, fn_traits, ptr_metadata, stmt_expr_attributes, entry_insert, map_try_insert, if_let_guard, bench_black_box, inline_const, inherent_associated_types, associated_type_bounds, let_chains, asm, destructuring_assignment)] 4 | //extern crate nom; 5 | extern crate thiserror; 6 | #[macro_use] 7 | extern crate enum_display_derive; 8 | extern crate yaxpeax_arch; 9 | extern crate yaxpeax_x86; 10 | extern crate goblin; 11 | extern crate target_lexicon; 12 | extern crate bitvec; 13 | extern crate bitflags; 14 | extern crate rangemap; 15 | 16 | use thiserror::Error; 17 | use std::fmt::Display; 18 | 19 | //mod parser; 20 | //use parser::parse_expr; 21 | //mod closure; 22 | //use closure::{CompileError, EvalError, compile_expr, generate_closure}; 23 | extern crate lineiform; 24 | use lineiform::{Lineiform}; 25 | use lineiform::tracer::{Tracer, TracerError}; 26 | use lineiform::block::{Function, BlockError}; 27 | use lineiform::lift::{Jit, LiftError, Location, JitValue}; 28 | 29 | use std::collections::HashMap; 30 | 31 | #[derive(Error, Debug)] 32 | pub enum MainError { 33 | //#[error("failed to compile: {0}")] 34 | //Compile(#[from] CompileError), 35 | //#[error("failed to evaluate: {0}")] 36 | //Eval(#[from] EvalError), 37 | #[error("tracer error: {0}")] 38 | Tracer(#[from] TracerError), 39 | #[error("block lifting error: {0}")] 40 | Block(#[from] BlockError), 41 | #[error("lifting error: {0}")] 42 | Lift(#[from] LiftError), 43 | } 44 | 45 | 46 | use core::hint::black_box; 47 | fn main() -> Result<(), MainError> { 48 | println!("Hello, world!"); 49 | let mut f = std::fs::File::open("/proc/self/maps").unwrap(); 50 | let out = std::io::stdout(); 51 | std::io::copy(&mut f, &mut out.lock()).unwrap(); 52 | 53 | { 54 | let mut jit = crate::Lineiform::new(); 55 | use core::hint::black_box; 56 | use core::num::Wrapping; 57 | let clos = jit.speedup(move |a: usize| { 58 | let val: usize; 59 | if a > 32 { 60 | val = (1 as usize); 61 | } else { 62 | val = (0 as usize); 63 | } 64 | val 65 | }); 66 | assert_eq!(clos(0), 0); 67 | assert_eq!(clos(31), 0); 68 | assert_eq!(clos(32), 0); 69 | assert_eq!(clos(33), 1); 70 | return Ok(()) 71 | } 72 | 73 | // let (s, sexpr) = parse_expr("(+ 1 (+ 2 3))") 74 | // .map_err(|e| MainError::Compile(CompileError::Parse(format!("{}", e).into())))?; 75 | // println!("{:?}", sexpr); 76 | // let clos = generate_closure(sexpr.clone())?; 77 | // let mut env = closure::Env { 78 | // variables: HashMap::new() 79 | // }; 80 | // println!("clos ptr: 0x{:?}", closure::get_ptr_from_closure(&clos)); 81 | // let clos2 = compile_expr("(* 2 6)")?; 82 | // println!("clos ptr: 0x{:?}", closure::get_ptr_from_closure(&clos2)); 83 | // println!("clos: {:?}", clos(&mut env)?); 84 | // println!("clos: {:?}", clos2(&mut env)?); 85 | 86 | pub trait TestFn = Fn(usize) -> usize; 87 | pub(crate) type TestExpr = Box< 88 | dyn TestFn + Send + Sync, 89 | >; 90 | use std::sync::Arc; 91 | use std::sync::Mutex; 92 | let external = Arc::new(Mutex::new(1)); 93 | let one: TestExpr = box move |e: usize| { 94 | if black_box(true) == true { 1 } else { 0 } 95 | }; 96 | let five: TestExpr = box move |e: usize| { 97 | let mut acc = e; 98 | for i in 0..5 { 99 | acc += one(e); 100 | } 101 | return acc; 102 | }; 103 | let eval = five(0); 104 | assert_eq!(eval, 5); 105 | let eval2 = five(1); 106 | assert_eq!(eval2, 6); 107 | 108 | use core::num::Wrapping; 109 | extern "C" fn test(a: Wrapping, b: Wrapping) -> Wrapping { 110 | a + Wrapping(2) 111 | } 112 | 113 | use yaxpeax_x86::long_mode::RegSpec; 114 | let mut tracer = Tracer::new()?; 115 | let assumed_test = Function::<(usize, usize), usize>::new(&mut tracer, test as *const ())? 116 | .assume_args(vec![(Location::Reg(RegSpec::rdi()), JitValue::Const(Wrapping(3)))]); 117 | let mut lifted = Jit::new(&mut tracer); 118 | let (lowered_test, size) = lifted.lower(assumed_test)?; 119 | unsafe { 120 | let jitted_pinned_test = std::mem::transmute::<_, extern "C" fn(usize, usize)->usize>(lowered_test); 121 | assert_eq!(jitted_pinned_test(1, 2), 5); 122 | } 123 | Ok(()) 124 | } 125 | 126 | fn test_jit() -> Result<(), MainError> { 127 | let mut tracer = Tracer::new()?; 128 | //use closure::ClosureFn; 129 | //let mut func = Function::from_fn(&mut tracer, &clos)?; 130 | use core::num::Wrapping; 131 | #[repr(C)] 132 | #[derive(Debug, PartialEq)] 133 | struct Oversized { 134 | a: Wrapping, 135 | b: Wrapping, 136 | c: Wrapping 137 | } 138 | extern "C" fn testcase(n: Wrapping) -> Oversized { 139 | Oversized { 140 | a: black_box(n + Wrapping(1)) + Wrapping(2), 141 | b: n, 142 | c: n + Wrapping(15) 143 | } 144 | } 145 | 146 | let mut func = Function::::new(&mut tracer, testcase as *const ())?; 147 | println!("base @ {:x}", func.base as usize); 148 | tracer.format(&func.instructions)?; 149 | 150 | let mut jitted = Jit::new(&mut tracer); 151 | let (new_func, new_size) = jitted.lower(func)?; 152 | let new_func_dis = tracer.disassemble(new_func as *const (), new_size as usize)?; 153 | println!("recompiled function:"); 154 | tracer.format(&new_func_dis)?; 155 | let callable: extern "C" fn(Wrapping) -> Oversized = 156 | unsafe { core::mem::transmute(new_func) }; 157 | assert_eq!(testcase(Wrapping(5)), callable(Wrapping(5))); 158 | 159 | Ok(()) 160 | } 161 | 162 | #[cfg(test)] 163 | mod test { 164 | use crate::*; 165 | #[test] 166 | pub fn test_can_pin_argument() -> Result<(), crate::MainError> { 167 | use core::num::Wrapping; 168 | extern "C" fn test(a: Wrapping, _b: Wrapping) -> Wrapping { 169 | a + Wrapping(2) 170 | } 171 | 172 | use yaxpeax_x86::long_mode::RegSpec; 173 | let mut tracer = Tracer::new()?; 174 | let assumed_test = Function::<(usize, usize), usize>::new(&mut tracer, test as *const ())? 175 | .assume_args(vec![ 176 | (Location::Reg(RegSpec::rdi()), JitValue::Const(Wrapping(3)))]); 177 | let mut lifted = Jit::new(&mut tracer); 178 | let (lowered_test, size) = lifted.lower(assumed_test)?; 179 | unsafe { 180 | let jitted_pinned_test = std::mem::transmute::<_, extern "C" fn(usize, usize)->usize>(lowered_test); 181 | assert_eq!(jitted_pinned_test(1, 2), 5); 182 | } 183 | Ok(()) 184 | } 185 | 186 | #[test] 187 | pub fn test_can_optimize_closure() -> Result<(), crate::MainError> { 188 | use core::num::Wrapping; 189 | let a: Wrapping = Wrapping(1); 190 | let mut jit = crate::Lineiform::new(); 191 | use core::hint::black_box; 192 | let clos = jit.speedup(move |()| { black_box(a) + Wrapping(2) }); 193 | assert_eq!(clos(()), Wrapping(3)); 194 | Ok(()) 195 | } 196 | 197 | #[test] 198 | pub fn test_can_optimize_two_closures() -> Result<(), crate::MainError> { 199 | use core::num::Wrapping; 200 | let a: Wrapping = Wrapping(1); 201 | let mut jit = crate::Lineiform::new(); 202 | use core::hint::black_box; 203 | let clos1 = black_box(move |()| { black_box(a) }); 204 | let clos2 = jit.speedup(move |()| { black_box(clos1)(()) + Wrapping(2) }); 205 | assert_eq!(clos2(()), Wrapping(3)); 206 | Ok(()) 207 | } 208 | 209 | #[test] 210 | pub fn test_passthrough_usize() -> Result<(), crate::MainError> { 211 | use core::num::Wrapping; 212 | let mut jit = crate::Lineiform::new(); 213 | use core::hint::black_box; 214 | let clos = jit.speedup(move |a: usize| { 215 | let b = black_box(0) + black_box(0); 216 | a + b 217 | }); 218 | for i in vec![0, !0, 0xFFFF_0000, 0x284591, &jit as &_ as *const _ as *const () as usize] { 219 | assert_eq!(clos(i), i); 220 | } 221 | Ok(()) 222 | } 223 | 224 | #[test] 225 | pub fn test_handles_wrapping_and_widths() -> Result<(), crate::MainError> { 226 | use core::num::Wrapping; 227 | let a: Wrapping = Wrapping(-3 as i32 as u32); 228 | let mut jit = crate::Lineiform::new(); 229 | use core::hint::black_box; 230 | let clos = jit.speedup(move |()| { black_box(a) + Wrapping(2) }); 231 | assert_eq!(clos(()), Wrapping(-1 as i32 as u32)); 232 | Ok(()) 233 | } 234 | 235 | 236 | #[test] 237 | pub fn test_handles_branches_const_true() -> Result<(), crate::MainError> { 238 | let a: u32 = 33; 239 | let mut jit = crate::Lineiform::new(); 240 | use core::hint::black_box; 241 | let clos = jit.speedup(move |()| { 242 | let val; 243 | if black_box(a) > 32 { 244 | val = 1; 245 | } else { 246 | val = 2; 247 | } 248 | val 249 | }); 250 | assert_eq!(clos(()), 1); 251 | Ok(()) 252 | } 253 | 254 | #[test] 255 | pub fn test_handles_branches_const_false() -> Result<(), crate::MainError> { 256 | let a: u32 = 32; 257 | let mut jit = crate::Lineiform::new(); 258 | use core::hint::black_box; 259 | let clos = jit.speedup(move |()| { 260 | let val; 261 | if black_box(a) > 32 { 262 | val = 1; 263 | } else { 264 | val = 2; 265 | } 266 | val 267 | }); 268 | assert_eq!(clos(()), 2); 269 | Ok(()) 270 | } 271 | 272 | #[test] 273 | pub fn test_handles_dyn_branches() -> Result<(), crate::MainError> { 274 | let mut jit = crate::Lineiform::new(); 275 | use core::hint::black_box; 276 | use core::num::Wrapping; 277 | let clos = jit.speedup(move |a: usize| { 278 | let val: usize; 279 | if a > 32 { 280 | val = (1 as usize); 281 | } else { 282 | val = (0 as usize); 283 | } 284 | val 285 | }); 286 | assert_eq!(clos(0), 0); 287 | assert_eq!(clos(31), 0); 288 | assert_eq!(clos(32), 0); 289 | assert_eq!(clos(33), 1); 290 | Ok(()) 291 | } 292 | 293 | #[test] 294 | pub fn test_handles_overflow_dyn_branches() -> Result<(), crate::MainError> { 295 | let mut jit = crate::Lineiform::new(); 296 | use core::hint::black_box; 297 | use core::num::Wrapping; 298 | let clos = jit.speedup(move |a: usize| { 299 | let val: usize; 300 | if Wrapping(a) + Wrapping(0x5) < Wrapping(a) { 301 | val = 1; 302 | } else { 303 | val = 2; 304 | } 305 | val 306 | }); 307 | assert_eq!(clos(0), 2); 308 | assert_eq!(clos((-1 as isize) as usize), 1); 309 | assert_eq!(clos((-6 as isize) as usize), 2); 310 | Ok(()) 311 | } 312 | 313 | #[test] 314 | pub fn test_handles_branches_dyn_blocks() -> Result<(), crate::MainError> { 315 | let mut jit = crate::Lineiform::new(); 316 | use core::hint::black_box; 317 | let clos = jit.speedup(move |a: usize| { 318 | let val: usize; 319 | if black_box(a) > 32 { 320 | val = 1; 321 | black_box(val); 322 | } else { 323 | val = 2; 324 | } 325 | val 326 | }); 327 | assert_eq!(clos(31), 2); 328 | assert_eq!(clos(32), 2); 329 | assert_eq!(clos(33), 1); 330 | Ok(()) 331 | } 332 | 333 | #[test] 334 | pub fn test_handles_multiple_branches() -> Result<(), crate::MainError> { 335 | let mut jit = crate::Lineiform::new(); 336 | use core::hint::black_box; 337 | let clos = jit.speedup(move |a: usize| { 338 | let mut val: usize; 339 | if black_box(a) > 32 { 340 | val = 1; 341 | black_box(val); 342 | } else { 343 | val = 2; 344 | } 345 | black_box(val); 346 | if black_box(a) > 64 { 347 | val += 1; 348 | black_box(val); 349 | } else { 350 | val += 2; 351 | } 352 | val 353 | }); 354 | assert_eq!(clos(31), 4); 355 | assert_eq!(clos(32), 4); 356 | assert_eq!(clos(33), 3); 357 | assert_eq!(clos(63), 3); 358 | assert_eq!(clos(64), 3); 359 | assert_eq!(clos(65), 2); 360 | Ok(()) 361 | } 362 | 363 | 364 | 365 | #[test] 366 | pub fn test_handles_loops() -> Result<(), crate::MainError> { 367 | use core::num::Wrapping; 368 | let a: Wrapping = Wrapping(10 as i32 as u32); 369 | let mut jit = crate::Lineiform::new(); 370 | use core::hint::black_box; 371 | let clos = jit.speedup(move |()| { 372 | let mut acc = Wrapping(0); 373 | for i in 0..black_box(a).0 { 374 | println!("i {} acc {}", i, acc); 375 | acc += acc + Wrapping(1); 376 | } 377 | acc 378 | }); 379 | assert_eq!(clos(()), Wrapping(1023 as i32 as u32)); 380 | Ok(()) 381 | } 382 | 383 | } 384 | -------------------------------------------------------------------------------- /src/parser.rs: -------------------------------------------------------------------------------- 1 | //! This example s-expr parser is taken from nom's examples (MIT license) 2 | //! In this example we build an [S-expression](https://en.wikipedia.org/wiki/S-expression) 3 | //! parser and tiny [lisp](https://en.wikipedia.org/wiki/Lisp_(programming_language)) interpreter. 4 | //! Lisp is a simple type of language made up of Atoms and Lists, forming easily parsable trees. 5 | 6 | #[global_allocator] 7 | static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc; 8 | 9 | use nom::{ 10 | branch::alt, 11 | bytes::complete::tag, 12 | character::complete::{alpha1, char, digit1, multispace0, multispace1, one_of}, 13 | combinator::{cut, map, map_res, opt}, 14 | error::{context, VerboseError}, 15 | multi::many0, 16 | sequence::{delimited, preceded, terminated, tuple}, 17 | IResult, Parser, 18 | }; 19 | use std::fmt::Display; 20 | 21 | /// We start by defining the types that define the shape of data that we want. 22 | /// In this case, we want something tree-like 23 | 24 | /// Starting from the most basic, we define some built-in functions that our lisp has 25 | #[derive(Debug, PartialEq, Clone, Copy, Display)] 26 | pub enum BuiltIn { 27 | Plus, 28 | Minus, 29 | Times, 30 | Divide, 31 | Equal, 32 | Not, 33 | Let, 34 | Set, 35 | Get, 36 | Do, 37 | Loop, 38 | } 39 | 40 | /// We now wrap this type and a few other primitives into our Atom type. 41 | /// Remember from before that Atoms form one half of our language. 42 | use std::num::Wrapping; 43 | pub type Int = Wrapping; 44 | 45 | #[derive(Debug, PartialEq, Clone, Display)] 46 | pub enum Atom { 47 | Num(Int), 48 | Keyword(String), 49 | //Boolean(bool), 50 | BuiltIn(BuiltIn), 51 | Unit, 52 | } 53 | 54 | /// The remaining half is Lists. We implement these as recursive Expressions. 55 | /// For a list of numbers, we have `'(1 2 3)`, which we'll parse to: 56 | /// ``` 57 | /// Expr::Quote(vec![Expr::Constant(Atom::Num(1)), 58 | /// Expr::Constant(Atom::Num(2)), 59 | /// Expr::Constant(Atom::Num(3))]) 60 | /// Quote takes an S-expression and prevents evaluation of it, making it a data 61 | /// structure that we can deal with programmatically. Thus any valid expression 62 | /// is also a valid data structure in Lisp itself. 63 | 64 | #[derive(Debug, PartialEq, Clone)] 65 | pub enum Expr { 66 | Constant(Atom), 67 | /// (func-name arg1 arg2) 68 | Application(Box, Vec), 69 | /// (if predicate do-this) 70 | If(Box, Box), 71 | /// (if predicate do-this otherwise-do-this) 72 | IfElse(Box, Box, Box), 73 | /// '(3 (if (+ 3 3) 4 5) 7) 74 | Quote(Vec), 75 | } 76 | impl Expr { 77 | pub fn need_int(&mut self) -> Result { 78 | match self { 79 | Expr::Constant(Atom::Num(u)) => Ok(*u), 80 | v => Err(crate::EvalError::TypeError("int", v._type())) 81 | } 82 | } 83 | 84 | pub fn need_keyword(&mut self) -> Result { 85 | match self { 86 | Expr::Constant(Atom::Keyword(u)) => Ok(u.to_string()), 87 | v => Err(crate::EvalError::TypeError("keyword", v._type())) 88 | } 89 | } 90 | 91 | pub fn _type(&self) -> &'static str { 92 | match self { 93 | Expr::Constant(Atom::Num(_)) => "int", 94 | Expr::Constant(Atom::BuiltIn(_)) => "builtin", 95 | _ => "unknown" 96 | } 97 | } 98 | } 99 | 100 | 101 | 102 | #[derive(Debug)] 103 | pub struct ListExpr(pub Vec); 104 | impl Display for ListExpr { 105 | fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> { 106 | let first = true; 107 | for elem in &self.0 { 108 | if first { 109 | write!(fmt, "{}", elem)?; 110 | } else { 111 | write!(fmt, " {}", elem)?; 112 | } 113 | } 114 | Ok(()) 115 | } 116 | } 117 | 118 | impl Display for Expr { 119 | fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> { 120 | match self { 121 | Expr::Constant(a) => write!(fmt, "{}", a), 122 | Expr::Application(func, body) => write!(fmt, "({} {})", func, ListExpr(body.clone())), 123 | Expr::If(cond, body) => write!(fmt, "(if {} {})", cond, body), 124 | Expr::IfElse(cond, t_body, f_body) => write!(fmt, "(if-else {} {} {})", 125 | cond, t_body, f_body), 126 | Expr::Quote(contents) => write!(fmt, "'({})", ListExpr(contents.clone())) 127 | } 128 | } 129 | } 130 | 131 | /// Continuing the trend of starting from the simplest piece and building up, 132 | /// we start by creating a parser for the built-in operator functions. 133 | fn parse_builtin_op<'a>(i: &'a str) -> IResult<&'a str, BuiltIn, VerboseError<&'a str>> { 134 | // one_of matches one of the characters we give it 135 | map(alt((tag("+"), tag("-"), tag("*"), tag("/"), tag("="), 136 | tag("let"), tag("set"), tag("get"), tag("do"), tag("loop"))), |stem: &str| 137 | // because we are matching single character tokens, we can do the matching logic 138 | // on the returned value 139 | match stem { 140 | "+" => BuiltIn::Plus, 141 | "-" => BuiltIn::Minus, 142 | "*" => BuiltIn::Times, 143 | "/" => BuiltIn::Divide, 144 | "=" => BuiltIn::Equal, 145 | "let" => BuiltIn::Let, 146 | "set" => BuiltIn::Set, 147 | "get" => BuiltIn::Get, 148 | "do" => BuiltIn::Do, 149 | "loop" => BuiltIn::Loop, 150 | _ => unreachable!(), 151 | })(i) 152 | } 153 | 154 | fn parse_builtin<'a>(i: &'a str) -> IResult<&'a str, BuiltIn, VerboseError<&'a str>> { 155 | // alt gives us the result of first parser that succeeds, of the series of 156 | // parsers we give it 157 | alt(( 158 | parse_builtin_op, 159 | // map lets us process the parsed output, in this case we know what we parsed, 160 | // so we ignore the input and return the BuiltIn directly 161 | map(tag("not"), |_| BuiltIn::Not), 162 | ))(i) 163 | } 164 | 165 | /// Our boolean values are also constant, so we can do it the same way 166 | fn parse_bool<'a>(i: &'a str) -> IResult<&'a str, Atom, VerboseError<&'a str>> { 167 | alt(( 168 | map(tag("#t"), |_| Atom::Num(Wrapping(1))), 169 | map(tag("#f"), |_| Atom::Num(Wrapping(0))), 170 | ))(i) 171 | } 172 | 173 | /// The next easiest thing to parse are keywords. 174 | /// We introduce some error handling combinators: `context` for human readable errors 175 | /// and `cut` to prevent back-tracking. 176 | /// 177 | /// Put plainly: `preceded(tag(":"), cut(alpha1))` means that once we see the `:` 178 | /// character, we have to see one or more alphabetic chararcters or the input is invalid. 179 | fn parse_keyword<'a>(i: &'a str) -> IResult<&'a str, Atom, VerboseError<&'a str>> { 180 | map( 181 | context("keyword", preceded(tag(":"), cut(alpha1))), 182 | |sym_str: &str| Atom::Keyword(sym_str.to_string()), 183 | )(i) 184 | } 185 | 186 | /// Next up is number parsing. We're keeping it simple here by accepting any number (> 1) 187 | /// of digits but ending the program if it doesn't fit into an i32. 188 | fn parse_num<'a>(i: &'a str) -> IResult<&'a str, Atom, VerboseError<&'a str>> { 189 | map_res(digit1, |digit_str: &str| { 190 | digit_str.parse::().map(Wrapping).map(Atom::Num) 191 | })(i) 192 | } 193 | 194 | /// Now we take all these simple parsers and connect them. 195 | /// We can now parse half of our language! 196 | fn parse_atom<'a>(i: &'a str) -> IResult<&'a str, Atom, VerboseError<&'a str>> { 197 | alt(( 198 | parse_num, 199 | parse_bool, 200 | map(parse_builtin, Atom::BuiltIn), 201 | parse_keyword, 202 | ))(i) 203 | } 204 | 205 | /// We then add the Expr layer on top 206 | fn parse_constant<'a>(i: &'a str) -> IResult<&'a str, Expr, VerboseError<&'a str>> { 207 | map(parse_atom, |atom| Expr::Constant(atom))(i) 208 | } 209 | 210 | /// Before continuing, we need a helper function to parse lists. 211 | /// A list starts with `(` and ends with a matching `)`. 212 | /// By putting whitespace and newline parsing here, we can avoid having to worry about it 213 | /// in much of the rest of the parser. 214 | /// 215 | /// Unlike the previous functions, this function doesn't take or consume input, instead it 216 | /// takes a parsing function and returns a new parsing function. 217 | fn s_exp<'a, O1, F>(inner: F) -> impl FnMut(&'a str) -> IResult<&'a str, O1, VerboseError<&'a str>> 218 | where 219 | F: Parser<&'a str, O1, VerboseError<&'a str>>, 220 | { 221 | delimited( 222 | char('('), 223 | preceded(multispace0, inner), 224 | context("closing paren", cut(preceded(multispace0, char(')')))), 225 | ) 226 | } 227 | 228 | /// We can now use our new combinator to define the rest of the `Expr`s. 229 | /// 230 | /// Starting with function application, we can see how the parser mirrors our data 231 | /// definitions: our definition is `Application(Box, Vec)`, so we know 232 | /// that we need to parse an expression and then parse 0 or more expressions, all 233 | /// wrapped in an S-expression. 234 | /// 235 | /// `tuple` is used to sequence parsers together, so we can translate this directly 236 | /// and then map over it to transform the output into an `Expr::Application` 237 | fn parse_application<'a>(i: &'a str) -> IResult<&'a str, Expr, VerboseError<&'a str>> { 238 | let application_inner = map(tuple((parse_expr, many0(parse_expr))), |(head, tail)| { 239 | Expr::Application(Box::new(head), tail) 240 | }); 241 | // finally, we wrap it in an s-expression 242 | s_exp(application_inner)(i) 243 | } 244 | 245 | /// Because `Expr::If` and `Expr::IfElse` are so similar (we easily could have 246 | /// defined `Expr::If` to have an `Option` for the else block), we parse both 247 | /// in a single function. 248 | /// 249 | /// In fact, we define our parser as if `Expr::If` was defined with an Option in it, 250 | /// we have the `opt` combinator which fits very nicely here. 251 | fn parse_if<'a>(i: &'a str) -> IResult<&'a str, Expr, VerboseError<&'a str>> { 252 | let if_inner = context( 253 | "if expression", 254 | map( 255 | preceded( 256 | // here to avoid ambiguity with other names starting with `if`, if we added 257 | // variables to our language, we say that if must be terminated by at least 258 | // one whitespace character 259 | terminated(tag("if"), multispace1), 260 | cut(tuple((parse_expr, parse_expr, opt(parse_expr)))), 261 | ), 262 | |(pred, true_branch, maybe_false_branch)| { 263 | if let Some(false_branch) = maybe_false_branch { 264 | Expr::IfElse( 265 | Box::new(pred), 266 | Box::new(true_branch), 267 | Box::new(false_branch), 268 | ) 269 | } else { 270 | Expr::If(Box::new(pred), Box::new(true_branch)) 271 | } 272 | }, 273 | ), 274 | ); 275 | s_exp(if_inner)(i) 276 | } 277 | 278 | /// A quoted S-expression is list data structure. 279 | /// 280 | /// This example doesn't have the symbol atom, but by adding variables and changing 281 | /// the definition of quote to not always be around an S-expression, we'd get them 282 | /// naturally. 283 | fn parse_quote<'a>(i: &'a str) -> IResult<&'a str, Expr, VerboseError<&'a str>> { 284 | // this should look very straight-forward after all we've done: 285 | // we find the `'` (quote) character, use cut to say that we're unambiguously 286 | // looking for an s-expression of 0 or more expressions, and then parse them 287 | map( 288 | context("quote", preceded(tag("'"), cut(s_exp(many0(parse_expr))))), 289 | |exprs| Expr::Quote(exprs), 290 | )(i) 291 | } 292 | 293 | /// We tie them all together again, making a top-level expression parser! 294 | 295 | pub fn parse_expr<'a>(i: &'a str) -> IResult<&'a str, Expr, VerboseError<&'a str>> { 296 | preceded( 297 | multispace0, 298 | alt((parse_constant, parse_application, parse_if, parse_quote)), 299 | )(i) 300 | } 301 | -------------------------------------------------------------------------------- /tangle/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "tangle" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | petgraph = "0.6.0" 10 | lazy_static = "1.4.0" 11 | yaxpeax-arch = "0.2.5" 12 | dynasm = "1.1.0" 13 | dynasmrt = "1.1.0" 14 | ena = "0.14.0" 15 | rangemap = "0.1.11" 16 | frunk = "0.4.0" 17 | ascent = "0.3.0" 18 | qcell = "0.5.2" 19 | fixedbitset = "*" 20 | 21 | [dependencies.yaxpeax-x86] 22 | version = "1.0.4" 23 | features = ["fmt"] 24 | path = "../../yaxpeax-x86" # FIXME when yax x86 1.1 comes out 25 | -------------------------------------------------------------------------------- /tangle/src/abi.rs: -------------------------------------------------------------------------------- 1 | use crate::region::Region; 2 | use crate::port::Storage; 3 | use yaxpeax_x86::long_mode::RegSpec; 4 | 5 | use core::fmt::Debug; 6 | 7 | pub trait Abi: Debug { 8 | fn provide_arguments(&self, request: Vec) -> Vec; 9 | fn provide_returns(&self, request: Vec) -> Vec; 10 | } 11 | 12 | #[derive(Default, Clone, Debug)] 13 | pub struct x86_64; 14 | 15 | pub enum AbiRequest { 16 | Integer(usize), 17 | Float(usize), 18 | Vector(usize), 19 | } 20 | 21 | #[derive(Clone, Debug)] 22 | pub enum AbiStorage { 23 | Register(RegSpec), 24 | // FPRegister 25 | // XMMRegister 26 | // etc 27 | StackSlot(usize), 28 | Problem(), 29 | } 30 | use AbiStorage::*; 31 | 32 | impl Abi for x86_64 { 33 | // this probably has to be something to do with a generic argument that has a Tracable trait 34 | fn provide_arguments(&self, request: Vec) -> Vec { 35 | request.iter().map(|req| { 36 | match req { 37 | AbiRequest::Integer(n) => 38 | vec![ 39 | Register(RegSpec::rdi()), 40 | Register(RegSpec::rsi()), 41 | Register(RegSpec::rdx()), 42 | Problem() 43 | ].drain(..).take(*n).collect::>(), 44 | _ => unimplemented!() 45 | } 46 | }).flatten().collect() 47 | } 48 | 49 | fn provide_returns(&self, request: Vec) -> Vec { 50 | request.iter().map(|req| { 51 | match req { 52 | AbiRequest::Integer(n) => 53 | vec![ 54 | Register(RegSpec::rax()), 55 | Register(RegSpec::rdx()), 56 | Problem() 57 | ].drain(..).take(*n).collect::>(), 58 | _ => unimplemented!() 59 | } 60 | }).flatten().collect() 61 | } 62 | } 63 | -------------------------------------------------------------------------------- /tangle/src/flags.rs: -------------------------------------------------------------------------------- 1 | #[derive(Hash, PartialEq, Clone, Copy, Display, Debug)] 2 | enum Flags { 3 | CF = 0, 4 | PF = 2, 5 | AF = 4, 6 | ZF = 6, 7 | SF = 7, 8 | TF = 8, 9 | IF = 9, 10 | DF = 10, 11 | OF = 11, 12 | FINAL = 12, 13 | } 14 | use Flags::*; 15 | 16 | /// A macro that allows for saving of `eflags` values to a Context after 17 | /// operations. Care should be taken that the last operation of `e` is the 18 | /// operation you want to collect flags from - you should probably only 19 | /// call this with some inline assembly expression to prevent the compiler 20 | /// from turning your a + b into a lea, for example. 21 | /// TODO: CPUID checks for if we have lahf 22 | macro_rules! getflags { 23 | ($ctx:expr, $e:expr, $bits:expr) => { 24 | $e; 25 | let mut flags: u64; 26 | asm!(" 27 | pushfq; 28 | pop {0}; 29 | ", out(reg) flags); 30 | println!("flags are 0b{:b}", flags); 31 | $ctx.set_flags_const($bits, flags); 32 | } 33 | } 34 | 35 | macro_rules! do_op { 36 | ($ctx:expr, $flag:expr, $left:expr, $right:expr, $asm:expr, $flags:expr, $op:tt) => { 37 | do_op!($ctx, $flag, reg, $left, $right, $asm, $flags, $op) 38 | }; 39 | ($ctx:expr, $flag:expr, $class:ident, $left:expr, $right:expr, $asm:expr, $flags:expr, $op:tt) => { 40 | if $flag { 41 | let mut val = $left; 42 | unsafe { 43 | getflags!($ctx, asm!($asm, inout($class) val, in($class) $right), 44 | $flags); 45 | } 46 | val 47 | } else { 48 | (Wrapping($left) $op (Wrapping($right))).0 49 | } 50 | } 51 | } 52 | 53 | /// A macro to create operations on all same-width JitTemp constants 54 | macro_rules! const_ops { 55 | ($ctx:expr, $flag:expr, $left:expr, $right:expr, $asm:expr, $flags:expr, $op:tt) => { { 56 | use JitTemp::*; 57 | match ($left, $right) { 58 | (Const8(l), Const8(r)) => 59 | Const8(do_op!($ctx, $flag, reg_byte, l, r, $asm, $flags, $op)), 60 | (Const16(l), Const16(r)) => 61 | Const16(do_op!($ctx, $flag, reg, l, r, $asm, $flags, $op)), 62 | (Const32(l), Const32(r)) => 63 | Const32(do_op!($ctx, $flag, reg, l, r, $asm, $flags, $op)), 64 | (Const64(l), Const64(r)) => 65 | Const64(do_op!($ctx, $flag, reg, l, r, $asm, $flags, $op)), 66 | _ => unimplemented!( 67 | "op {} for {:?} {:?}", 68 | stringify!($op), $left, $right), 69 | } 70 | } } 71 | } 72 | 73 | -------------------------------------------------------------------------------- /tangle/src/ir.rs: -------------------------------------------------------------------------------- 1 | // Cranelift (and probably any other codegen IR like LLVM or libgccjit) is just 2 | // a bit too abstact for us; specifically, they target multiple ISAs, so try not 3 | // to expose processor flags behavior to the user. 4 | // We *need* processor flag behavior, however, because the assembly we're lifting 5 | // depends on it. 6 | // 7 | // Thus, this IR was born. It's basically just a thin shim over dynasm-rs, but 8 | // with an RSVDG-esque graph structure for data dependencies for instructions. 9 | // This is because the *primary* use for it is register allocation; the assembly 10 | // itself is already reasonably optimized just by LLVM when it is first emitted, 11 | // we just want to elide checks and stop everything from moving back and forth for 12 | // the function ABI for each inlined closure body. 13 | 14 | use petgraph::stable_graph::StableGraph; 15 | use petgraph::graph::{NodeIndex}; 16 | use yaxpeax_x86::long_mode::{RegSpec}; 17 | use dynasmrt::x64::Assembler; 18 | use dynasmrt::{AssemblyOffset, DynasmApi}; 19 | 20 | use std::collections::{HashMap, HashSet}; 21 | 22 | use crate::node::{NodeBehavior, NodeVariant, NodeIdx, NodeOwner}; 23 | pub use crate::port::{PortIdx, EdgeVariant}; 24 | use crate::region::{Region, RegionIdx, RegionEdge}; 25 | use crate::abi::Abi; 26 | //use crate::opt::PassRunner; 27 | use crate::select::PatternManager; 28 | use crate::node::Continuation; 29 | 30 | // yaxpeax decoder example 31 | mod decoder { 32 | use yaxpeax_arch::{Decoder, Reader, ReaderBuilder}; 33 | 34 | pub fn decode_stream< 35 | 'data, 36 | A: yaxpeax_arch::Arch, 37 | U: ReaderBuilder, 38 | >(data: U) -> Vec 39 | where A::Instruction: std::fmt::Display 40 | { 41 | let mut reader = ReaderBuilder::read_from(data); 42 | let mut address: A::Address = reader.total_offset(); 43 | 44 | let decoder = A::Decoder::default(); 45 | let mut decode_res = decoder.decode(&mut reader); 46 | let mut res = Vec::new(); 47 | loop { 48 | match decode_res { 49 | Ok(inst) => { 50 | //println!("{}: {}", address.show(), inst); 51 | decode_res = decoder.decode(&mut reader); 52 | address = reader.total_offset(); 53 | res.push(inst); 54 | } 55 | Err(e) => { 56 | //println!("{}: decode error: {}", address.show(), e); 57 | break; 58 | } 59 | } 60 | } 61 | res 62 | } 63 | } 64 | 65 | 66 | #[derive(Clone)] 67 | pub struct VirtualRegister { 68 | pub ports: Vec, // these are reversed: the 0th element is the last port that uses it 69 | pub hints: HashSet, 70 | pub backing: Option, 71 | pub allocated: bool, 72 | } 73 | pub type VirtualRegisterMap = HashMap; 74 | 75 | #[derive(Default)] 76 | pub struct IR { 77 | body: Option, 78 | /// A directed acyclic graph of regions; if there is an edge r_1 -> r_2, then 79 | /// r_2 is within r_1. 80 | regions: StableGraph, 81 | /// The outer-most region that all other regions are under. 82 | master_region: RegionIdx, 83 | pub owner: Option, 84 | } 85 | 86 | impl IR { 87 | pub fn new() -> Self { 88 | let mut r = Region::new(); 89 | r.live = true; 90 | let mut r_map = StableGraph::new(); 91 | let r_x = r_map.add_node(r); 92 | Self { 93 | body: None, 94 | regions: r_map, 95 | master_region: r_x, 96 | owner: Some(NodeOwner::new()), 97 | } 98 | } 99 | 100 | /// Create a region in this IR instance at the top-most level. 101 | pub fn new_region(&mut self, abi: Option) -> RegionIdx { 102 | let mut r = Region::new(); 103 | r.live = true; 104 | r.abi = abi.map(|a| (box a) as Box); 105 | let r_x = self.regions.add_node(r); 106 | { self.regions.node_weight_mut(r_x).unwrap().idx = r_x; } 107 | self.regions.add_edge(self.master_region, r_x, ()); 108 | r_x 109 | } 110 | 111 | pub fn add_function(&mut self, f: NodeVariant::Function) { 112 | let mut owner = self.owner.take().unwrap(); 113 | self.in_region(self.master_region, |r, ir| { 114 | r.add_node(&mut owner, f, |n, r| { 115 | }); 116 | }); 117 | self.owner = Some(owner); 118 | } 119 | 120 | pub fn set_body(&mut self, f: NodeVariant::Function) { 121 | println!("setting IR body"); 122 | let body = self.in_region(self.master_region, |r, ir| { 123 | NodeIndex::new(r.nodes.node_count()) 124 | }); 125 | assert_eq!(f.args as usize, A); 126 | if f.cont == Continuation::Return { 127 | assert_eq!(f.outs as usize, O); 128 | } 129 | self.add_function(f); 130 | self.body = Some(body); // once told me 131 | } 132 | 133 | pub fn in_region(&mut self, r: RegionIdx, f: F) -> O 134 | where F: FnOnce(&mut Region, &mut Self) -> O { 135 | let mut dummy = Region::new(); 136 | std::mem::swap(&mut dummy, &mut self.regions[r]); 137 | assert_eq!(dummy.live, true, "dead region? {:?}", r); 138 | //assert!(dummy.idx.is_some(), "region {:?} not connected to graph", r); 139 | let o = f(&mut dummy, self); 140 | std::mem::swap(&mut dummy, &mut self.regions[r]); 141 | o 142 | } 143 | 144 | pub fn optimize(&mut self) -> &mut Self { 145 | self 146 | } 147 | 148 | pub fn regalloc(&mut self) { 149 | for r in self.regions.node_weights_mut() { 150 | if r.idx == self.master_region { 151 | continue; 152 | } 153 | r.apply_constraints(); 154 | r.attach_ports(); 155 | r.propogate_state_edges(); 156 | r.observe_state_outputs(); 157 | let mut virt_map: VirtualRegisterMap = HashMap::new(); 158 | r.connect_block_params(&mut virt_map, self.owner.as_ref().unwrap()); 159 | println!("propogated state edges"); 160 | 161 | //r.move_constants_to_operands(); 162 | r.remove_nops(self.owner.as_mut().unwrap()); 163 | // Create virtual register storages for all nodes and ports 164 | r.create_virtual_registers(&mut virt_map, self.owner.as_mut().unwrap()); 165 | r.create_dependencies(); // we have to re-add node dependencies for the inserted movs 166 | 167 | let mut patterner = PatternManager::default(); 168 | patterner.run(self.owner.as_ref().unwrap(), r, &mut virt_map); 169 | //let runner = PassRunner; 170 | //runner.run(r); 171 | //panic!(); 172 | 173 | println!("created virtual registers"); 174 | r.annotate_port_times_and_hints(&mut virt_map); 175 | println!("annoated port times and hints"); 176 | r.allocate_physical_for_virtual(&mut virt_map); 177 | println!("allocated physical registers for virtual register"); 178 | //r.replace_virtual_with_backing(&mut virt_map); 179 | //println!("replacing virtual registers with physical registers"); 180 | // Create a dependency graph for registers 181 | //let dg = r.create_register_dependencies(); 182 | patterner.virt_map = virt_map; 183 | r.patterns = Some(patterner); 184 | } 185 | } 186 | 187 | pub fn validate(&mut self) { 188 | // This is currently worse than useless 189 | //for r in self.regions.node_weights_mut() { 190 | // if r.idx == self.master_region { 191 | // continue; 192 | // } 193 | // r.validate(); 194 | //} 195 | } 196 | 197 | pub fn codegen(&mut self) -> (Assembler, AssemblyOffset, usize) { 198 | let mut ops = Assembler::new().unwrap(); 199 | let start = ops.offset(); 200 | let mut ret = None; 201 | let mut owner = self.owner.take().unwrap(); 202 | self.in_region(self.master_region, |r, ir| { 203 | let mut n = crate::region::NodeGraph::default(); 204 | std::mem::swap(&mut r.nodes, &mut n); 205 | ret = Some(n[ir.body.unwrap()].codegen(&owner, vec![], vec![], r, ir, &mut ops)); 206 | std::mem::swap(&mut r.nodes, &mut n); 207 | }); 208 | let end = ops.offset(); 209 | (ops, start, end.0 - start.0) 210 | } 211 | 212 | pub fn compile_fn(&mut self) -> Result<(extern "C" fn(A) -> O, usize), Box> { 213 | self.regalloc(); 214 | self.validate(); 215 | let (mut ops, off, size) = self.codegen(); 216 | let buf = ops.finalize().unwrap(); 217 | let hello_fn: extern "C" fn(A) -> O = unsafe { std::mem::transmute(buf.ptr(off)) }; 218 | let all = crate::ir::decoder::decode_stream::(unsafe { 219 | core::slice::from_raw_parts( 220 | buf.ptr(off) as *const u8, 221 | size as usize) as &[u8] 222 | }); 223 | println!("disassembly: \n"); 224 | let mut fmt = String::new(); 225 | for inst in all { 226 | inst.write_to(&mut fmt)?; 227 | fmt.push('\n'); 228 | } 229 | println!("{}", fmt); 230 | std::mem::forget(buf); 231 | Ok((hello_fn, size)) 232 | } 233 | 234 | pub fn print(&self) { 235 | // XXX: we want to dump a graphviz view or something else pretty here; 236 | // debugging a graph based IR without that sounds like hell. 237 | } 238 | } 239 | 240 | #[cfg(test)] 241 | mod test { 242 | use super::*; 243 | use crate::node::{Node, Operation}; 244 | 245 | type S = NodeVariant::Simple; 246 | 247 | #[test] 248 | pub fn simple_function() { 249 | let mut ir = IR::new(); 250 | let mut f = Node::function::<0, 0>(&mut ir); 251 | } 252 | 253 | #[test] 254 | pub fn function_inc() { 255 | let mut ir = IR::new(); 256 | let mut f = Node::function::<1, 1>(&mut ir); 257 | let port = f.add_argument(&mut ir); 258 | let ret_port = f.add_return(&mut ir); 259 | let mut inc = Node::simple(Operation::Inc); 260 | f.add_body(inc, &mut ir, |inc, ir| { 261 | inc.connect_input(0, port, ir); 262 | inc.connect_output(0, ret_port, ir); 263 | }); 264 | ir.set_body(f); 265 | } 266 | 267 | #[test] 268 | pub fn function_inc_regalloc() { 269 | let mut ir = IR::new(); 270 | let mut f = Node::function::<1, 1>(&mut ir); 271 | let port = f.add_argument(&mut ir); 272 | let ret_port = f.add_return(&mut ir); 273 | let mut inc = Node::simple(Operation::Inc); 274 | f.add_body(inc, &mut ir, |inc, r| { 275 | inc.connect_input(0, port, r); 276 | inc.connect_output(0, ret_port, r); 277 | }); 278 | ir.set_body(f); 279 | ir.regalloc(); 280 | ir.validate(); 281 | } 282 | 283 | #[test] 284 | pub fn function_inc_codegen() { 285 | let mut ir = IR::new(); 286 | let mut f = Node::function::<1, 1>(&mut ir); 287 | let port = f.add_argument(&mut ir); 288 | let ret_port = f.add_return(&mut ir); 289 | let mut inc = Node::simple(Operation::Inc); 290 | f.add_body(inc, &mut ir, |inc, r| { 291 | inc.connect_input(0, port, r); 292 | inc.connect_output(0, ret_port, r); 293 | }); 294 | ir.set_body(f); 295 | let hello_fn: extern "C" fn(usize) -> usize = ir.compile_fn().unwrap().0; 296 | assert_eq!(hello_fn(1), 2); 297 | } 298 | 299 | #[test] 300 | pub fn function_add_codegen() { 301 | let mut ir = IR::new(); 302 | let mut f = Node::function::<2, 1>(&mut ir); 303 | let port_0 = f.add_argument(&mut ir); 304 | let port_1 = f.add_argument(&mut ir); 305 | let ret_port = f.add_return(&mut ir); 306 | let mut add = Node::simple(Operation::Add); 307 | f.add_body(add, &mut ir, |add, r| { 308 | add.connect_input(0, port_0, r); 309 | add.connect_input(1, port_1, r); 310 | add.connect_output(0, ret_port, r); 311 | }); 312 | ir.set_body(f); 313 | let hello_fn: extern "C" fn((usize, usize)) -> usize = ir.compile_fn().unwrap().0; 314 | assert_eq!(hello_fn((1, 2)), 3); 315 | } 316 | 317 | #[test] 318 | pub fn function_inc_disjoint() { 319 | let mut ir = IR::new(); 320 | let mut f = Node::function::<2, 2>(&mut ir); 321 | let port_0 = f.add_argument(&mut ir); 322 | let port_1 = f.add_argument(&mut ir); 323 | let ret_port_0 = f.add_return(&mut ir); 324 | let ret_port_1 = f.add_return(&mut ir); 325 | let mut inc_0 = Node::simple(Operation::Inc); 326 | let mut inc_1 = Node::simple(Operation::Inc); 327 | f.add_body(inc_0, &mut ir, |inc, r| { 328 | inc.connect_input(0, port_0, r); 329 | inc.connect_output(0, ret_port_0, r); 330 | }); 331 | f.add_body(inc_1, &mut ir, |inc, r| { 332 | inc.connect_input(0, port_1, r); 333 | inc.connect_output(0, ret_port_1, r); 334 | }); 335 | 336 | ir.set_body(f); 337 | let hello_fn: extern "C" fn((usize, usize)) -> (usize, usize) = ir.compile_fn().unwrap().0; 338 | assert_eq!(hello_fn((1, 2)), (2, 3)); 339 | } 340 | 341 | #[test] 342 | pub fn function_inc_then_add() { 343 | let mut ir = IR::new(); 344 | let mut f = Node::function::<2, 1>(&mut ir); 345 | let port_0 = f.add_argument(&mut ir); 346 | let port_1 = f.add_argument(&mut ir); 347 | let ret_port = f.add_return(&mut ir); 348 | let mut inc_0 = Node::simple(Operation::Inc); 349 | let mut inc_1 = Node::simple(Operation::Inc); 350 | let mut add = Node::simple(Operation::Add); 351 | let inc_0_i = f.add_body(inc_0, &mut ir, |inc_0, r| { 352 | inc_0.connect_input(0, port_0, r); 353 | }).0; 354 | let inc_1_i = f.add_body(inc_1, &mut ir, |inc_1, r| { 355 | inc_1.connect_input(0, port_1, r); 356 | }).0; 357 | f.add_body(add, &mut ir, |add, r| { 358 | add.connect_input(0, r.nodes[inc_0_i].sinks()[0], r); 359 | add.connect_input(1, r.nodes[inc_1_i].sinks()[0], r); 360 | add.connect_output(0, ret_port, r); 361 | }); 362 | ir.set_body(f); 363 | 364 | let hello_fn: extern "C" fn((usize, usize)) -> usize = ir.compile_fn().unwrap().0; 365 | assert_eq!(hello_fn((1, 2)), 5); 366 | } 367 | 368 | #[test] 369 | pub fn function_inc_shared() { 370 | let mut ir = IR::new(); 371 | let mut f = Node::function::<2, 2>(&mut ir); 372 | let port_0 = f.add_argument(&mut ir); 373 | let port_1 = f.add_argument(&mut ir); 374 | let ret_port_0 = f.add_return(&mut ir); 375 | let ret_port_1 = f.add_return(&mut ir); 376 | let mut inc_0 = Node::simple(Operation::Inc); 377 | let mut inc_1 = Node::simple(Operation::Inc); 378 | f.add_body(inc_0, &mut ir, |inc, r| { 379 | inc.connect_input(0, port_0, r); 380 | inc.connect_output(0, ret_port_0, r); 381 | }); 382 | f.add_body(inc_1, &mut ir, |inc, r| { 383 | inc.connect_input(0, port_1, r); 384 | inc.connect_output(0, ret_port_1, r); 385 | }); 386 | 387 | ir.set_body(f); 388 | let hello_fn: extern "C" fn((usize, usize)) -> (usize, usize) = ir.compile_fn().unwrap().0; 389 | assert_eq!(hello_fn((1, 2)), (2, 3)); 390 | } 391 | 392 | #[test] 393 | pub fn function_inc_a_lot() { 394 | let mut ir = IR::new(); 395 | let mut f = Node::function::<1, 1>(&mut ir); 396 | let mut input = f.add_argument(&mut ir); 397 | let output = f.add_return(&mut ir); 398 | // TODO: make the port virtual register proprogation better so we can 399 | // bump this up 400 | let count = 200; 401 | for i in 0..count { 402 | let mut inc_0 = Node::simple(Operation::Inc); 403 | f.add_body(inc_0, &mut ir, |inc, r| { 404 | inc.connect_input(0, input, r); 405 | input = inc.sinks()[0]; 406 | }); 407 | } 408 | ir.in_region(f.region, |r, ir| { 409 | r.connect_ports(input, output); 410 | }); 411 | 412 | 413 | ir.set_body(f); 414 | let hello_fn: extern "C" fn(usize) -> usize = ir.compile_fn().unwrap().0; 415 | assert_eq!(hello_fn(0), count); 416 | } 417 | 418 | #[test] 419 | pub fn function_add_constant() { 420 | let mut ir = IR::new(); 421 | let mut f = Node::function::<1, 1>(&mut ir); 422 | let input = f.add_argument(&mut ir); 423 | let output = f.add_return(&mut ir); 424 | 425 | let mut two = Node::constant(2); 426 | let mut add = Node::simple(Operation::Add); 427 | let mut two_const = None; 428 | f.add_body(two, &mut ir, |two, r| { 429 | two_const = Some(two.sinks()[0]); 430 | }); 431 | f.add_body(add, &mut ir, |add, r| { 432 | add.connect_input(0, input, r); 433 | add.connect_input(1, two_const.unwrap(), r); 434 | add.connect_output(0, output, r); 435 | }); 436 | 437 | ir.set_body(f); 438 | let hello_fn: extern "C" fn(usize) -> usize = ir.compile_fn().unwrap().0; 439 | assert_eq!(hello_fn(0), 2); 440 | assert_eq!(hello_fn(1), 3); 441 | } 442 | 443 | #[test] 444 | pub fn function_bb_simple() { 445 | let mut ir = IR::new(); 446 | let mut f = Node::function::<1, 1>(&mut ir); 447 | let input = f.add_argument(&mut ir); 448 | let output = f.add_return(&mut ir); 449 | 450 | // entry[input]: br_call bb1, input; 451 | // bb1[bp0]: add bp0, 2; mov output, bp0; 452 | 453 | let mut bb = Node::bb(); 454 | let (bb0, bp_0) = f.add_body(bb, &mut ir, |bb, mut r| { 455 | // create basic block parameter 456 | (bb.sinks()[0], bb.add_output(&mut r)) 457 | }).1; 458 | 459 | let mut two = Node::constant(2); 460 | let mut add = Node::simple(Operation::Add); 461 | let mut two_const = None; 462 | f.add_body(two, &mut ir, |two, r| { 463 | two_const = Some(two.sinks()[0]); 464 | }); 465 | f.add_body(add, &mut ir, |add, r| { 466 | add.connect_input(0, bp_0, r); 467 | add.connect_input(1, two_const.unwrap(), r); 468 | add.connect_output(0, output, r); 469 | }); 470 | 471 | let mut bcall = Node::bcall(); 472 | f.add_body(bcall, &mut ir, |bcall, mut r| { 473 | // call bb0 474 | bcall.connect_input(0, bb0, r); 475 | // with arg0 = input 476 | let arg = r.add_port(); 477 | r.connect_ports(input, arg); 478 | bcall.add_input(arg, &mut r); 479 | }); 480 | 481 | ir.set_body(f); 482 | let hello_fn: extern "C" fn(usize) -> usize = ir.compile_fn().unwrap().0; 483 | assert_eq!(hello_fn(0), 2); 484 | assert_eq!(hello_fn(1), 3); 485 | } 486 | 487 | #[test] 488 | pub fn function_bb_if_always_0() { 489 | let mut ir = IR::new(); 490 | let mut f = Node::function::<1, 1>(&mut ir); 491 | let input = f.add_argument(&mut ir); 492 | let output = f.add_return(&mut ir); 493 | 494 | // entry[input]: zero = 0; br_if zero, bb0, bb1, input; 495 | // bb0[bp0]: add bp0, 1; mov output, bp0; 496 | // bb1[bp0]: add bp0, 2; mov output, bp0; 497 | 498 | let mut bb = Node::bb(); 499 | let (bb0, bp_0) = f.add_body(bb, &mut ir, |bb, mut r| { 500 | // create basic block parameter 501 | (bb.sinks()[0], bb.add_output(&mut r)) 502 | }).1; 503 | let mut bb = Node::bb(); 504 | let (bb1, bp_1) = f.add_body(bb, &mut ir, |bb, mut r| { 505 | // create basic block parameter 506 | (bb.sinks()[0], bb.add_output(&mut r)) 507 | }).1; 508 | 509 | let mut one = f.add_body(Node::constant(1), &mut ir, |n, r| { n.sinks()[0] }).1; 510 | let mut add = Node::simple(Operation::Add); 511 | f.add_body(add, &mut ir, |add, r| { 512 | add.connect_input(0, bp_0, r); 513 | add.connect_input(1, one, r); 514 | add.connect_output(0, output, r); 515 | }); 516 | 517 | let mut two = f.add_body(Node::constant(2), &mut ir, |n, r| { n.sinks()[0] }).1; 518 | let mut add = Node::simple(Operation::Add); 519 | f.add_body(add, &mut ir, |add, r| { 520 | add.connect_input(0, bp_1, r); 521 | add.connect_input(1, two, r); 522 | add.connect_output(0, output, r); 523 | }); 524 | 525 | let mut zero = f.add_body(Node::constant(0), &mut ir, |n, r| { n.sinks()[0] }).1; 526 | 527 | let mut bif = Node::bif(); 528 | f.add_body(bif, &mut ir, |bif, mut r| { 529 | // br_if zero, bb0, bb1 530 | bif.connect_input(0, zero, r); 531 | bif.connect_input(1, bb0, r); 532 | bif.connect_input(2, bb1, r); 533 | // with arg0 = input 534 | let arg = r.add_port(); 535 | r.connect_ports(input, arg); 536 | bif.add_input(arg, &mut r); 537 | }); 538 | 539 | ir.set_body(f); 540 | let hello_fn: extern "C" fn(usize) -> usize = ir.compile_fn().unwrap().0; 541 | assert_eq!(hello_fn(0), 1); 542 | assert_eq!(hello_fn(1), 2); 543 | } 544 | 545 | #[test] 546 | pub fn function_bb_if_always_1() { 547 | let mut ir = IR::new(); 548 | let mut f = Node::function::<1, 1>(&mut ir); 549 | let input = f.add_argument(&mut ir); 550 | let output = f.add_return(&mut ir); 551 | 552 | // entry[input]: zero = 0; br_if zero, bb0, bb1, input; 553 | // bb0[bp0]: add bp0, 1; call term, bp0; 554 | // bb1[bp0]: add bp0, 2; call term, bp0; 555 | // term[output]: ret output; 556 | let term = f.add_body(Node::bb(), &mut ir, |bb, mut r| { 557 | let arg = bb.add_output(r); 558 | r.connect_ports(arg, output); 559 | bb.sinks()[0] 560 | }).1; 561 | 562 | let (bb0, bp_0) = f.add_body(Node::bb(), &mut ir, |bb, mut r| { 563 | // create basic block parameter 564 | (bb.sinks()[0], bb.add_output(&mut r)) 565 | }).1; 566 | let (bb1, bp_1) = f.add_body(Node::bb(), &mut ir, |bb, mut r| { 567 | // create basic block parameter 568 | (bb.sinks()[0], bb.add_output(&mut r)) 569 | }).1; 570 | 571 | let mut one = f.add_body(Node::constant(1), &mut ir, |n, r| { n.sinks()[0] }).1; 572 | let add_0 = f.add_body(Node::simple(Operation::Add), &mut ir, |add, r| { 573 | add.connect_input(0, bp_0, r); 574 | add.connect_input(1, one, r); 575 | add.sinks()[0] 576 | }).1; 577 | f.add_body(Node::bcall(), &mut ir, |call, r| { 578 | call.connect_input(0, term, r); 579 | let param = r.add_port(); 580 | r.connect_ports(add_0, param); 581 | call.add_input(param, r); 582 | }); 583 | 584 | let mut two = f.add_body(Node::constant(2), &mut ir, |n, r| { n.sinks()[0] }).1; 585 | let add_1 = f.add_body(Node::simple(Operation::Add), &mut ir, |add, r| { 586 | add.connect_input(0, bp_1, r); 587 | add.connect_input(1, two, r); 588 | add.sinks()[0] 589 | }).1; 590 | f.add_body(Node::bcall(), &mut ir, |call, r| { 591 | call.connect_input(0, term, r); 592 | let param = r.add_port(); 593 | r.connect_ports(add_1, param); 594 | call.add_input(param, r); 595 | }); 596 | 597 | let mut one = f.add_body(Node::constant(1), &mut ir, |n, r| { n.sinks()[0] }).1; 598 | 599 | let mut bif = Node::bif(); 600 | f.add_body(bif, &mut ir, |bif, mut r| { 601 | // br_if zero, bb0, bb1 602 | bif.connect_input(0, one, r); 603 | bif.connect_input(1, bb0, r); 604 | bif.connect_input(2, bb1, r); 605 | // with arg0 = input 606 | let arg = r.add_port(); 607 | r.connect_ports(input, arg); 608 | bif.add_input(arg, &mut r); 609 | }); 610 | 611 | ir.set_body(f); 612 | let hello_fn: extern "C" fn(usize) -> usize = ir.compile_fn().unwrap().0; 613 | assert_eq!(hello_fn(0), 2); 614 | assert_eq!(hello_fn(1), 3); 615 | } 616 | 617 | #[test] 618 | pub fn function_use_stack() { 619 | let mut ir = IR::new(); 620 | let mut f = Node::function::<1, 1>(&mut ir); 621 | let input = f.add_argument(&mut ir); 622 | let output = f.add_return(&mut ir); 623 | // We create a stack slot 624 | let mut ss1 = f.add_stack_slot(&mut ir); 625 | 626 | // ...and push, which should resolve to a use of ss1 627 | let mut push = Node::simple(Operation::StoreStack); 628 | ss1 = f.add_body(push, &mut ir, |i, r| { 629 | i.connect_input(0, ss1, r); 630 | i.connect_input(1, input, r); 631 | i.sinks()[0] 632 | }).1; 633 | 634 | // get the value 635 | let mut pop = Node::simple(Operation::LoadStack); 636 | let ss_val = f.add_body(pop, &mut ir, |i, r| { 637 | i.connect_input(0, ss1, r); 638 | i.sinks()[0] 639 | }).1; 640 | 641 | // then add to it via ss1 642 | let mut two = Node::constant(2); 643 | let mut add = Node::simple(Operation::Add); 644 | let mut two_const = f.add_body(two, &mut ir, |two, r| { 645 | two.sinks()[0] 646 | }).1; 647 | let added_val = f.add_body(add, &mut ir, |add, r| { 648 | println!("add closure {}", r.ports[ss1].storage); 649 | add.connect_input(0, ss_val, r); 650 | add.connect_input(1, two_const, r); 651 | add.sinks()[0] 652 | }).1; 653 | 654 | // and store it again 655 | let mut push = Node::simple(Operation::StoreStack); 656 | f.add_body(push, &mut ir, |i, r| { 657 | i.connect_input(0, ss1, r); 658 | i.connect_input(1, added_val, r); 659 | ss1 = i.sinks()[0]; 660 | }); 661 | 662 | // ...then pop it into the result register 663 | let mut pop = Node::simple(Operation::LoadStack); 664 | f.add_body(pop, &mut ir, |i, r| { 665 | i.connect_input(0, ss1, r); 666 | i.connect_output(0, output, r); 667 | }); 668 | 669 | ir.set_body(f); 670 | let hello_fn: extern "C" fn(usize) -> usize = ir.compile_fn().unwrap().0; 671 | assert_eq!(hello_fn(0), 2); 672 | assert_eq!(hello_fn(1), 3); 673 | } 674 | 675 | #[test] 676 | pub fn function_stack_ordering() { 677 | let mut ir = IR::new(); 678 | let mut f = Node::function::<1, 1>(&mut ir); 679 | let input = f.add_argument(&mut ir); 680 | let output = f.add_return(&mut ir); 681 | // We create a stack slot 682 | let mut ss1 = f.add_stack_slot(&mut ir); 683 | let initial_ss = ss1; 684 | 685 | // ...and push, which should resolve to a use of ss1 686 | let mut push = Node::simple(Operation::StoreStack); 687 | ss1 = f.add_body(push, &mut ir, |i, r| { 688 | i.connect_input(0, ss1, r); 689 | i.connect_input(1, input, r); 690 | i.sinks()[0] 691 | }).1; 692 | 693 | let mut pop = Node::simple(Operation::LoadStack); 694 | let ss_val_0 = f.add_body(pop, &mut ir, |i, r| { 695 | i.connect_input(0, ss1, r); 696 | i.sinks()[0] 697 | }).1; 698 | 699 | // then add to it via ss1 700 | let mut two = Node::constant(2); 701 | let mut add = Node::simple(Operation::Add); 702 | let mut two_const = f.add_body(two, &mut ir, |two, r| { 703 | two.sinks()[0] 704 | }).1; 705 | let added_val = f.add_body(add, &mut ir, |add, r| { 706 | println!("add closure {}", r.ports[ss1].storage); 707 | add.connect_input(0, ss_val_0, r); 708 | add.connect_input(1, two_const, r); 709 | add.sinks()[0] 710 | }).1; 711 | 712 | // and store it again 713 | let mut push = Node::simple(Operation::StoreStack); 714 | f.add_body(push, &mut ir, |i, r| { 715 | i.connect_input(0, ss1, r); 716 | i.connect_input(1, added_val, r); 717 | ss1 = i.sinks()[0]; 718 | }); 719 | 720 | // make a use of the *initial* stack value from before the load 721 | let mut pop = Node::simple(Operation::LoadStack); 722 | let ss_val_1 = f.add_body(pop, &mut ir, |i, r| { 723 | i.connect_input(0, initial_ss, r); 724 | i.sinks()[0] 725 | }).1; 726 | 727 | let mut three = Node::constant(3); 728 | let mut add = Node::simple(Operation::Add); 729 | let mut three_const = f.add_body(three, &mut ir, |three, r| { 730 | three.sinks()[0] 731 | }).1; 732 | f.add_body(add, &mut ir, |add, r| { 733 | add.connect_input(0, ss_val_1, r); 734 | add.connect_input(1, three_const, r); 735 | add.connect_output(0, output, r); 736 | }); 737 | 738 | 739 | ir.set_body(f); 740 | let hello_fn: extern "C" fn(usize) -> usize = ir.compile_fn().unwrap().0; 741 | assert_eq!(hello_fn(0), 3); 742 | assert_eq!(hello_fn(1), 4); 743 | } 744 | 745 | #[test] 746 | pub fn function_tailcall_external() { 747 | let mut ir = IR::new(); 748 | let mut f = Node::function::<1, 1>(&mut ir); 749 | let input = f.add_argument(&mut ir); 750 | let output = f.add_return(&mut ir); 751 | 752 | extern "C" fn external_function(a: usize) -> usize { 753 | a + 2 754 | } 755 | 756 | let mut addr = Node::constant(external_function as *const () as isize); 757 | let mut tailcall = Node::leave(); 758 | let mut addr_const = None; 759 | f.add_body(addr, &mut ir, |addr, r| { 760 | addr_const = Some(addr.sinks()[0]); 761 | }); 762 | f.add_body(tailcall, &mut ir, |tailcall, r| { 763 | tailcall.connect_input(0, addr_const.unwrap(), r); 764 | }); 765 | 766 | ir.set_body(f); 767 | let hello_fn: extern "C" fn(usize) -> usize = ir.compile_fn().unwrap().0; 768 | assert_eq!(hello_fn(0), 2); 769 | assert_eq!(hello_fn(1), 3); 770 | } 771 | 772 | #[test] 773 | pub fn function_tailcall_with_shuffle() { 774 | let mut ir = IR::new(); 775 | let mut f = Node::function::<2, 1>(&mut ir); 776 | let input_0 = f.add_argument(&mut ir); 777 | let input_1 = f.add_argument(&mut ir); 778 | let output = f.add_return(&mut ir); 779 | 780 | extern "C" fn external_function(a: usize) -> usize { 781 | a + 2 782 | } 783 | 784 | let mut addr = Node::constant(external_function as *const () as isize); 785 | let mut tailcall = Node::leave(); 786 | let mut addr_const = None; 787 | f.add_body(addr, &mut ir, |addr, r| { 788 | addr_const = Some(addr.sinks()[0]); 789 | }); 790 | f.add_body(tailcall, &mut ir, |tailcall, r| { 791 | tailcall.connect_input(0, addr_const.unwrap(), r); 792 | let arg = r.add_port(); 793 | r.constrain(arg, crate::port::Storage::Physical(RegSpec::rdi())); 794 | tailcall.add_input(arg, r); 795 | tailcall.connect_input(1, input_1, r); 796 | }); 797 | 798 | ir.set_body(f); 799 | let hello_fn: extern "C" fn((usize, usize)) -> usize = ir.compile_fn().unwrap().0; 800 | assert_eq!(hello_fn((0, 0)), 2); 801 | assert_eq!(hello_fn((0, 1)), 3); 802 | } 803 | } 804 | -------------------------------------------------------------------------------- /tangle/src/lib.rs: -------------------------------------------------------------------------------- 1 | #![allow(unused_variables, unused_mut, unused_imports, non_snake_case, dead_code, non_camel_case_types)] 2 | #![feature(box_syntax, let_chains, trait_upcasting, associated_type_defaults)] 3 | #![deny(unused_must_use, improper_ctypes_definitions)] 4 | extern crate lazy_static; 5 | extern crate petgraph; 6 | extern crate yaxpeax_x86; 7 | extern crate dynasmrt; 8 | extern crate ena; 9 | extern crate rangemap; 10 | extern crate fixedbitset; 11 | pub mod port; 12 | pub mod node; 13 | pub mod region; 14 | pub mod ir; 15 | pub mod time; 16 | pub mod abi; 17 | pub mod opt; 18 | pub mod select; 19 | pub use ir::*; 20 | -------------------------------------------------------------------------------- /tangle/src/node.rs: -------------------------------------------------------------------------------- 1 | use dynasmrt::x64::Assembler; 2 | use dynasmrt::{dynasm, DynasmApi}; 3 | use yaxpeax_x86::long_mode::RegSpec; 4 | use yaxpeax_x86::long_mode::{register_class, Instruction}; 5 | use petgraph::graph::{NodeIndex, EdgeIndex}; 6 | 7 | use crate::time::Timestamp; 8 | use crate::ir::{IR}; 9 | use crate::port::{PortMeta, PortIdx, PortEdge, Storage, EdgeVariant}; 10 | use crate::region::{Region, RegionIdx, State, StateVariant}; 11 | use crate::abi::{x86_64, Abi}; 12 | 13 | use std::rc::Rc; 14 | #[derive(Hash, PartialOrd, Ord, PartialEq, Eq, Clone, Copy, Default)] 15 | pub struct NodeIdxToken(u32); 16 | pub type NodeIdx = NodeIndex; 17 | pub type NodeEdge = EdgeIndex; 18 | unsafe impl petgraph::matrix_graph::IndexType for NodeIdxToken { 19 | #[inline(always)] 20 | fn new(x: usize) -> Self { 21 | NodeIdxToken(x as u32) 22 | } 23 | #[inline(always)] 24 | fn index(&self) -> usize { 25 | self.0 as usize 26 | } 27 | #[inline(always)] 28 | fn max() -> Self { 29 | NodeIdxToken(::std::u32::MAX) 30 | } 31 | } 32 | 33 | impl core::fmt::Debug for NodeIdxToken { 34 | fn fmt(&self, fmt: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> { 35 | write!(fmt, "n{}", self.0) 36 | } 37 | } 38 | 39 | use qcell::{TCell, TCellOwner}; 40 | use qcell::{QCell, QCellOwner}; 41 | //pub struct NodeMarker; 42 | pub type NodeOwner = QCellOwner; 43 | pub type NodeCell = QCell>; 44 | 45 | //#[derive(Debug)] 46 | pub struct Node { 47 | pub variant: Rc, 48 | pub inputs: Vec, 49 | pub outputs: Vec, 50 | pub label: Option, 51 | pub containing_region: Option, 52 | pub time: Timestamp, 53 | pub done: bool, 54 | } 55 | 56 | #[derive(Debug, PartialEq, Clone)] 57 | pub enum Operation { 58 | Nop, 59 | Apply, // Call a Function node with arguments 60 | Inc, 61 | Add, 62 | LoadStack, 63 | StoreStack 64 | } 65 | 66 | #[derive(Debug, PartialEq)] 67 | pub enum Continuation { 68 | Return, 69 | Jmp(PortIdx), 70 | } 71 | 72 | pub mod NodeVariant { 73 | use super::{Operation, Region}; 74 | 75 | use super::*; 76 | #[derive(Debug, Clone)] 77 | pub struct Constant(pub isize); 78 | #[derive(Debug, Clone)] 79 | pub struct Move(pub Storage, pub Storage); // A move operation, sink <- source. 80 | #[derive(Debug, Clone)] 81 | pub struct Simple(pub Operation); // Instructions or constant operands 82 | #[derive(Debug)] 83 | pub struct Function { 84 | pub args: u8, 85 | pub outs: u8, 86 | pub stack_slots: Vec, 87 | pub highwater: u8, 88 | pub region: RegionIdx, 89 | pub cont: Continuation, 90 | } // "Lambda-Nodes"; procedures and functions 91 | impl Function { 92 | pub fn new(ir: &mut IR) -> Function { 93 | let r = ir.new_region(Some(::default())); 94 | Self { 95 | region: r, 96 | args: 0, 97 | outs: 0, 98 | stack_slots: vec![], 99 | highwater: 0, 100 | cont: Continuation::Return, 101 | } 102 | } 103 | 104 | pub fn add_body<'id, T, F, F_O>( 105 | &mut self, 106 | mut n: T, 107 | ir: &mut IR, 108 | f: F) -> (NodeIdx, F_O) where F: FnOnce(&mut Node, &mut Region) -> F_O, T: NodeBehavior + core::any::Any + 'static + core::fmt::Debug 109 | { 110 | ir.in_region(self.region, |mut r, ir| { 111 | r.add_node(ir.owner.as_mut().unwrap(), n, f) 112 | }) 113 | } 114 | 115 | pub fn add_argument(&mut self, ir: &mut IR) -> PortIdx { 116 | let port = ir.in_region(self.region, |mut r, ir| { 117 | let port = r.add_source(); 118 | port 119 | }); 120 | self.args += 1; 121 | port 122 | } 123 | 124 | pub fn add_return(&mut self, ir: &mut IR) -> PortIdx { 125 | let port = ir.in_region(self.region, |mut r, ir| { 126 | let port = r.add_sink(); 127 | port 128 | }); 129 | self.outs += 1; 130 | port 131 | } 132 | 133 | pub fn add_stack_slot(&mut self, ir: &mut IR) -> PortIdx { 134 | let port = ir.in_region(self.region, |mut r, ir| { 135 | let new_state = r.states.len(); 136 | let port = r.add_port(); 137 | r.states.push(State { 138 | name: "stack".to_string(), 139 | variant: StateVariant::Stack(self.stack_slots.len().try_into().unwrap()), 140 | producers: vec![], 141 | }); 142 | 143 | r.constrain(port, Storage::Immaterial(Some(new_state.try_into().unwrap()))); 144 | r.ports[port].set_variant(EdgeVariant::State); 145 | port 146 | }); 147 | 148 | self.stack_slots.push(port); 149 | self.highwater = core::cmp::max(self.highwater, self.stack_slots.len().try_into().unwrap()); 150 | port 151 | } 152 | } 153 | #[derive(Debug)] 154 | pub struct BrEntry; // "Block Entry"; header of a basic block, source of bb params 155 | #[derive(Debug)] 156 | pub struct BrCall; // "Block Call"; jump to a block entry, with arguments 157 | #[derive(Debug)] 158 | pub struct BrIf; // "Branch If"; jump to either the first or second block's entry, with arguments, depending on the switch value. 159 | /// Tailcalls or jumps into native code are represented as Leave nodes. 160 | /// These nodes should have: 161 | /// 1) an input port for the address to jump to 162 | /// 2) some arbitrary number of input state or value ports for values 163 | /// that should be observed by the jump target (e.g. register state or side-effects) 164 | /// 3) an output port state edge that should be connected to the function return 165 | #[derive(Debug)] 166 | pub struct Leave; 167 | } 168 | 169 | pub use NodeVariant::*; 170 | 171 | pub trait NodeBehavior: core::fmt::Debug + core::any::Any { 172 | fn set_time(&mut self, time: Timestamp) { 173 | unimplemented!(); 174 | } 175 | fn get_time(&self) -> Timestamp { 176 | unimplemented!(); 177 | } 178 | 179 | fn create_ports(&self, r: &mut Region) { 180 | } 181 | 182 | fn ports_callback(&mut self, inputs: Vec, outputs: Vec, r: &mut Region) { 183 | } 184 | 185 | fn codegen(&self, token: &NodeOwner, inputs: Vec, outputs: Vec, r: &mut Region, ir: &mut IR, ops: &mut Assembler) { 186 | unimplemented!(); 187 | } 188 | 189 | 190 | fn tag(&self) -> String; 191 | 192 | fn input_count(&self) -> usize { 193 | unimplemented!(); 194 | } 195 | fn output_count(&self) -> usize { 196 | unimplemented!(); 197 | } 198 | } 199 | 200 | impl Node { 201 | pub fn as_variant<'a, T>(&'a self, token: &'a NodeOwner) -> Option<&'a T> where T: 'static { 202 | (&*self.variant.ro(token).deref() as &dyn NodeBehavior as &(dyn core::any::Any + 'static)).downcast_ref::() 203 | } 204 | 205 | pub fn as_variant_mut<'a, T>(&'a mut self, token: &'a mut NodeOwner) -> Option<&'a mut T> where T: 'static + Clone { 206 | let x = self.as_variant::(token); 207 | //(&mut x?.clone() as &mut (dyn core::any::Any + 'static)).downcast_mut::() 208 | panic!() 209 | } 210 | 211 | fn set_time(&mut self, time: Timestamp) { 212 | self.time = time; 213 | } 214 | fn get_time(&self) -> Timestamp { 215 | self.time 216 | } 217 | 218 | pub fn sinks(&self) -> Vec { 219 | self.outputs.clone() 220 | } 221 | 222 | pub fn sources(&self) -> Vec { 223 | self.inputs.clone() 224 | } 225 | 226 | pub fn ports(&self) -> Vec { 227 | let mut p = self.sources(); 228 | p.append(&mut self.sinks()); 229 | p 230 | } 231 | 232 | pub fn add_input(&mut self, p: PortIdx, r: &mut Region) { 233 | self.inputs.push(p) 234 | } 235 | 236 | pub fn add_output(&mut self, r: &mut Region) -> PortIdx { 237 | let p_x = r.add_port(); 238 | self.outputs.push(p_x); 239 | p_x 240 | } 241 | 242 | pub fn connect_input<'a>(&'a mut self, idx: usize, input: PortIdx, r: &'a mut Region) -> PortEdge { 243 | let p = self.inputs[idx]; 244 | r.connect_ports(input, p) 245 | } 246 | 247 | pub fn connect_output(&mut self, idx: usize, output: PortIdx, r: &mut Region) -> PortEdge { 248 | let p = self.outputs[idx]; 249 | r.connect_ports(p, output) 250 | } 251 | 252 | pub fn connect_operands(&mut self, input: usize, output: usize, r: &mut Region) -> PortEdge { 253 | let input = self.inputs[input]; 254 | let output = self.outputs[output]; 255 | r.connect_ports(input, output) 256 | } 257 | 258 | pub fn create_ports(&mut self, token: &mut NodeOwner, r: &mut Region) { 259 | println!("create_ports called"); 260 | self.variant.rw(token).create_ports(r); 261 | let mut inputs = vec![]; 262 | let mut outputs = vec![]; 263 | for i in 0..self.input_count(token) { 264 | let p = r.add_port(); 265 | self.add_input(p, r); 266 | inputs.push(p); 267 | } 268 | for i in 0..self.output_count(token) { 269 | outputs.push(self.add_output(r)); 270 | } 271 | self.variant.rw(token).ports_callback(inputs, outputs, r); 272 | } 273 | 274 | pub fn input_count(&self, token: &NodeOwner) -> usize { 275 | self.variant.ro(token).input_count() 276 | } 277 | 278 | pub fn output_count(&self, token: &NodeOwner) -> usize { 279 | self.variant.ro(token).output_count() 280 | } 281 | 282 | pub fn codegen(&self, token: &NodeOwner, inputs: Vec, outputs: Vec, r: &mut Region, ir: &mut IR, ops: &mut Assembler) { 283 | println!("node codegen for {:?} @ {}", self.variant.ro(token), self.time); 284 | self.variant.ro(token).codegen(token, self.inputs.clone(), self.outputs.clone(), r, ir, ops) 285 | } 286 | } 287 | 288 | impl NodeBehavior for NodeVariant::Constant { 289 | fn input_count(&self) -> usize { 290 | 0 291 | } 292 | 293 | fn output_count(&self) -> usize { 294 | 1 295 | } 296 | 297 | fn codegen(&self, token: &NodeOwner, inputs: Vec, outputs: Vec, r: &mut Region, ir: &mut IR, ops: &mut Assembler) { 298 | let output = &r.ports[outputs[0]]; 299 | match output.storage.as_ref().unwrap() { 300 | Storage::Physical(r0) => match r0.class() { 301 | register_class::Q => dynasm!(ops 302 | ; mov Rq(r0.num()), QWORD (self.0) as i64 303 | ), 304 | x => unimplemented!("unknown class {:?} for constant", x), 305 | }, 306 | _ => unimplemented!(), 307 | } 308 | } 309 | 310 | fn tag(&self) -> String { 311 | self.0.to_string() 312 | } 313 | } 314 | 315 | impl NodeBehavior for NodeVariant::BrEntry { 316 | fn input_count(&self) -> usize { 317 | 0 318 | } 319 | 320 | fn output_count(&self) -> usize { 321 | // Start with one port output, which is the block itself for use in br_call 322 | 1 323 | } 324 | 325 | fn ports_callback(&mut self, inputs: Vec, outputs: Vec, r: &mut Region) { 326 | // set the block port as a state edge, of the block itself 327 | let port = outputs[0]; 328 | let new_state = r.states.len(); 329 | r.states.push(State { 330 | name: format!("bb{}", new_state), 331 | variant: StateVariant::Block(NodeIdx::new(0)), // this gets filled in by Region::observe_state_edges 332 | producers: vec![], 333 | }); 334 | 335 | r.constrain(port, Storage::Immaterial(Some(new_state.try_into().unwrap()))); 336 | r.ports[port].set_variant(EdgeVariant::State); 337 | } 338 | 339 | fn codegen(&self, token: &NodeOwner, inputs: Vec, outputs: Vec, r: &mut Region, ir: &mut IR, ops: &mut Assembler) { 340 | // block entries don't codegen to anything 341 | } 342 | 343 | fn tag(&self) -> String { 344 | "br_entry".to_string() 345 | } 346 | } 347 | 348 | impl NodeBehavior for NodeVariant::BrCall { 349 | fn input_count(&self) -> usize { 350 | // takes as its first input the block to jump to 351 | 1 352 | } 353 | 354 | fn output_count(&self) -> usize { 355 | 0 356 | } 357 | 358 | fn ports_callback(&mut self, inputs: Vec, outputs: Vec, r: &mut Region) { 359 | r.ports[inputs[0]].set_variant(EdgeVariant::State); 360 | } 361 | 362 | fn codegen(&self, token: &NodeOwner, inputs: Vec, outputs: Vec, r: &mut Region, ir: &mut IR, ops: &mut Assembler) { 363 | panic!() 364 | } 365 | 366 | fn tag(&self) -> String { 367 | "br_call".to_string() 368 | } 369 | } 370 | 371 | impl NodeBehavior for NodeVariant::BrIf { 372 | fn input_count(&self) -> usize { 373 | // takes as its first input the switch value, and the second and third as the 0 or 1 switch 374 | // case. 375 | 3 376 | } 377 | 378 | fn output_count(&self) -> usize { 379 | 0 380 | } 381 | 382 | fn ports_callback(&mut self, inputs: Vec, outputs: Vec, r: &mut Region) { 383 | r.ports[inputs[1]].set_variant(EdgeVariant::State); 384 | r.ports[inputs[2]].set_variant(EdgeVariant::State); 385 | } 386 | 387 | fn codegen(&self, token: &NodeOwner, inputs: Vec, outputs: Vec, r: &mut Region, ir: &mut IR, ops: &mut Assembler) { 388 | panic!() 389 | } 390 | 391 | fn tag(&self) -> String { 392 | "br_if".to_string() 393 | } 394 | } 395 | 396 | 397 | 398 | impl NodeBehavior for NodeVariant::Simple { 399 | fn input_count(&self) -> usize { 400 | match &self.0 { 401 | Operation::Inc => 1, 402 | Operation::Add => 2, 403 | Operation::LoadStack => 1, 404 | Operation::StoreStack => 2, 405 | Operation::Nop => 0, 406 | Operation::Apply => unimplemented!() 407 | } 408 | } 409 | 410 | fn output_count(&self) -> usize { 411 | match &self.0 { 412 | Operation::Inc => 1, 413 | Operation::Add => 1, 414 | Operation::LoadStack | Operation::StoreStack => 1, 415 | Operation::Nop => 0, 416 | Operation::Apply => unimplemented!() 417 | } 418 | } 419 | 420 | 421 | fn codegen(&self, token: &NodeOwner, inputs: Vec, outputs: Vec, r: &mut Region, ir: &mut IR, ops: &mut Assembler) { 422 | match &self.0 { 423 | Operation::Inc => { 424 | let operand = &r.ports[inputs[0]]; 425 | match operand.storage.as_ref().unwrap() { 426 | Storage::Physical(r) => match r.class() { 427 | register_class::Q => dynasm!(ops 428 | ; inc Rq(r.num()) 429 | ), 430 | register_class::D => dynasm!(ops 431 | ; inc Rd(r.num()) 432 | ), 433 | x => unimplemented!("unknown class {:?} for inc", x), 434 | }, 435 | _ => unimplemented!(), 436 | } 437 | }, 438 | Operation::Add => { 439 | let operand_0 = &r.ports[inputs[0]]; 440 | let operand_1 = &r.ports[inputs[1]]; 441 | match (operand_0.storage.as_ref().unwrap(), operand_1.storage.as_ref().unwrap()) { 442 | (Storage::Physical(r0), Storage::Physical(r1)) => match (r0.class(), r1.class()) { 443 | (register_class::Q, register_class::Q) => dynasm!(ops 444 | ; add Rq(r0.num()), Rq(r1.num()) 445 | ), 446 | (register_class::D, register_class::D) => dynasm!(ops 447 | ; add Rd(r0.num()), Rd(r1.num()) 448 | ), 449 | x => unimplemented!("unknown class {:?} for add", x), 450 | }, 451 | //(Storage::Physical(r0), Storage::Immaterial(None)) => match r0.class() { 452 | // register_class::Q => dynasm!(ops 453 | // ; add Rq(r0.num()), 454 | // operand_1.get_meta::().unwrap().0.try_into().unwrap() 455 | // ), 456 | // x => unimplemented!("unknown class {:?} for add", x), 457 | //} 458 | _ => unimplemented!(), 459 | } 460 | }, 461 | // val <- ss 462 | // out: val' 463 | // THIS IS FUCKED: 464 | // stack slots actually have to be a STATE edge and not a data edge, 465 | // so that we can propogate them without having any conflicts. 466 | // otherwise we have a Stack storage for the ss, but not any uses, 467 | // and physical registers propogate upwards into it and it doesn't 468 | // know what it is anymore. 469 | Operation::LoadStack => { 470 | let ss = &r.ports[inputs[0]]; 471 | let output = &r.ports[outputs[0]]; 472 | NodeVariant::Move::codegen(&NodeVariant::Move(output.storage.unwrap(), ss.storage.unwrap()), 473 | token, vec![inputs[0]], vec![outputs[0]], r, ir, ops); 474 | //match output.storage.as_ref().unwrap() { 475 | // Storage::Physical(r0) => match r0.class() { 476 | // register_class::Q => dynasm!(ops 477 | // ; mov Rq(r0.num()), QWORD [rsp+ss_off] 478 | // ), 479 | // x => unimplemented!("unknown class {:?} for load_ss", x) 480 | // }, 481 | // _ => unimplemented!(), 482 | //} 483 | }, 484 | // ss <- val 485 | // out: ss' 486 | Operation::StoreStack => { 487 | let ss = &r.ports[inputs[0]]; 488 | let input = &r.ports[inputs[1]]; 489 | if let Storage::Immaterial(Some(state)) = ss.storage.unwrap() { 490 | if let Some(state) = r.states.get(state as usize) { 491 | NodeVariant::Move::codegen(&NodeVariant::Move(ss.storage.unwrap(), input.storage.unwrap()), 492 | token, vec![inputs[1]], vec![inputs[0]], r, ir, ops); 493 | } else { 494 | panic!("bad state for ss storage") 495 | } 496 | } else { panic!("bad storage for ss {:?}", ss) } 497 | //let ss_off = ss.get_meta::().unwrap().0; 498 | //match input.storage.as_ref().unwrap() { 499 | // Storage::Physical(r0) => match r0.class() { 500 | // register_class::Q => dynasm!(ops 501 | // ; mov QWORD [rsp+ss_off], Rq(r0.num()) 502 | // ), 503 | // x => unimplemented!("unknown class {:?} for load_ss", x) 504 | // }, 505 | // _ => unimplemented!(), 506 | //} 507 | }, 508 | x => unimplemented!("unimplemented codegen for {:?}", x), 509 | } 510 | } 511 | 512 | fn ports_callback(&mut self, inputs: Vec, outputs: Vec, r: &mut Region) { 513 | println!("{:?} <- {:?} ", outputs, inputs); 514 | match &self.0 { 515 | Operation::Inc => 516 | { r.connect_ports(inputs[0], outputs[0]); }, 517 | Operation::Add => 518 | { r.connect_ports(inputs[0], outputs[0]); }, 519 | Operation::LoadStack => { r.ports[inputs[0]].set_variant(EdgeVariant::State); }, 520 | Operation::StoreStack => { 521 | r.ports[inputs[0]].set_variant(EdgeVariant::State); 522 | r.ports[outputs[0]].set_variant(EdgeVariant::State); 523 | r.connect_ports(inputs[0], outputs[0]); 524 | }, // this connects the state ports 525 | Operation::Nop => {}, 526 | Operation::Apply => unimplemented!(), 527 | }; 528 | } 529 | 530 | fn tag(&self) -> String { 531 | match &self.0 { 532 | Operation::Inc => "inc".to_string(), 533 | Operation::Add => "add".to_string(), 534 | Operation::LoadStack => "load_ss".to_string(), 535 | Operation::StoreStack => "store_ss".to_string(), 536 | Operation::Nop => "nop".to_string(), 537 | Operation::Apply => "apply".to_string(), 538 | } 539 | } 540 | } 541 | 542 | impl NodeBehavior for NodeVariant::Leave { 543 | fn tag(&self) -> String { 544 | "leave".to_string() 545 | } 546 | 547 | fn input_count(&self) -> usize { 548 | 1 549 | } 550 | 551 | fn output_count(&self) -> usize { 552 | 1 553 | } 554 | 555 | fn ports_callback(&mut self, inputs: Vec, outputs: Vec, r: &mut Region) { 556 | println!("tailcall leave "); 557 | r.constrain(outputs[0], Storage::Immaterial(None)); 558 | } 559 | 560 | fn codegen(&self, token: &NodeOwner, inputs: Vec, outputs: Vec, r: &mut Region, ir: &mut IR, ops: &mut Assembler) { 561 | // "what happens if you try to tailcall with stack slots allocated?" 562 | // haha. yeah. 563 | let source = &r.ports[inputs[0]]; 564 | match source.storage.as_ref().unwrap() { 565 | Storage::Physical(r0) => { 566 | match r0.class() { 567 | register_class::Q => { 568 | dynasm!(ops 569 | ; jmp Rq(r0.num()) 570 | ); 571 | }, 572 | _ => unimplemented!("tailcall class {:?}", r0.class()) 573 | } 574 | }, 575 | _ => unimplemented!("unknown tailcall target") 576 | } 577 | } 578 | } 579 | 580 | impl NodeBehavior for NodeVariant::Move { 581 | 582 | fn tag(&self) -> String { 583 | format!("mov {} <- {}", self.0, self.1) 584 | } 585 | 586 | fn input_count(&self) -> usize { 587 | 1 588 | } 589 | 590 | fn output_count(&self) -> usize { 591 | 1 592 | } 593 | 594 | fn ports_callback(&mut self, inputs: Vec, outputs: Vec, r: &mut Region) { 595 | println!("{:?} <- {:?} ", outputs, inputs); 596 | r.constrain(outputs[0], self.0.clone()); 597 | r.constrain(inputs[0], self.1.clone()); 598 | } 599 | 600 | fn codegen(&self, token: &NodeOwner, inputs: Vec, outputs: Vec, r: &mut Region, ir: &mut IR, ops: &mut Assembler) { 601 | let sink = &r.ports[outputs[0]]; 602 | let source = &r.ports[inputs[0]]; 603 | match (sink.storage.as_ref().unwrap(), source.storage.as_ref().unwrap()) { 604 | (Storage::Physical(r0), Storage::Physical(r1)) => { 605 | // we can simply not emit any move at all if it was trying to move 606 | // a register into itself 607 | if r0 == r1 { return; } 608 | match (r0.class(), r1.class()) { 609 | (register_class::Q, register_class::Q) => 610 | dynasm!(ops 611 | ; mov Rq(r0.num()), Rq(r1.num()) 612 | ), 613 | (x, y) => unimplemented!("{:?} <- {:?} unimplemented", x, y), 614 | } 615 | }, 616 | (Storage::Physical(r0), Storage::Immaterial(Some(state))) => { 617 | let off = match r.states[*state as usize].variant { 618 | StateVariant::Stack(o) => (o * 8).try_into().unwrap(), 619 | _ => panic!(), 620 | }; 621 | match r0.class() { 622 | register_class::Q => dynasm!(ops 623 | ; mov Rq(r0.num()), QWORD [rsp+off] 624 | ), 625 | _ => unimplemented!() 626 | } 627 | }, 628 | (Storage::Immaterial(Some(state)), Storage::Physical(r1)) => { 629 | let off = match r.states[*state as usize].variant { 630 | StateVariant::Stack(o) => (o * 8).try_into().unwrap(), 631 | _ => panic!(), 632 | }; 633 | match r1.class() { 634 | register_class::Q => dynasm!(ops 635 | ; mov QWORD [rsp+off], Rq(r1.num()) 636 | ), 637 | _ => unimplemented!() 638 | } 639 | }, 640 | (x, y) => unimplemented!("codegen for {:?} <- {:?}", x, y), 641 | } 642 | } 643 | 644 | } 645 | 646 | 647 | impl NodeBehavior for NodeVariant::Function { 648 | fn tag(&self) -> String { 649 | "function".to_string() 650 | } 651 | 652 | fn input_count(&self) -> usize { 653 | 0 654 | } 655 | 656 | fn output_count(&self) -> usize { 657 | // The output of a function node is always the function itself 658 | // todo lol 659 | 1 660 | } 661 | 662 | fn codegen(&self, token: &NodeOwner, inputs: Vec, outputs: Vec, r: &mut Region, ir: &mut IR, ops: &mut Assembler) { 663 | assert_eq!(self.args as usize, A); 664 | if self.highwater != 0 { 665 | // allocate stack space 666 | dynasm!(ops 667 | ; sub rsp, (self.highwater*8).try_into().unwrap() 668 | ); 669 | } 670 | // if we're calling codegen on a function, it should be the only one. 671 | ir.in_region(self.region, |r, ir| { r.codegen(token, ir, ops); }); 672 | if self.highwater != 0 { 673 | // cleanup stack space 674 | dynasm!(ops 675 | ; add rsp, (self.highwater*8).try_into().unwrap() 676 | ); 677 | } 678 | match self.cont { 679 | Continuation::Return => { 680 | assert_eq!(self.outs as usize, O); 681 | dynasm!(ops 682 | ; ret 683 | ) 684 | }, 685 | Continuation::Jmp(o) => { 686 | // TODO: assert that the o port is connected to a Leave node 687 | }, 688 | } 689 | } 690 | } 691 | 692 | use core::ops::Deref; 693 | impl Deref for Node { 694 | type Target = Rc; 695 | fn deref(&self) -> &::Target { 696 | &self.variant 697 | } 698 | } 699 | use core::ops::DerefMut; 700 | impl DerefMut for Node { 701 | fn deref_mut(&mut self) -> &mut ::Target { 702 | &mut self.variant 703 | } 704 | } 705 | 706 | impl Node { 707 | pub fn new(owner: &NodeOwner, var: Box) -> Node { 708 | Node { 709 | variant: Rc::new(NodeCell::new(owner, var)), 710 | inputs: vec![], 711 | outputs: vec![], 712 | label: None, 713 | containing_region: None, 714 | time: Timestamp::new().increment(), 715 | done: false, 716 | } 717 | } 718 | pub fn constant(n: isize) -> NodeVariant::Constant { 719 | NodeVariant::Constant(n) 720 | } 721 | 722 | pub fn simple(op: Operation) -> NodeVariant::Simple { 723 | NodeVariant::Simple(op) 724 | } 725 | 726 | pub fn r#move(sink: Storage, source: Storage) -> NodeVariant::Move { 727 | NodeVariant::Move(sink, source) 728 | } 729 | 730 | // id like for this to just take an ABI parameter but rust deprececated 731 | // default generic arguments for functions :/ 732 | pub fn function(ir: &mut IR) -> NodeVariant::Function { 733 | NodeVariant::Function::new::(ir) 734 | } 735 | 736 | pub fn leave() -> NodeVariant::Leave { 737 | NodeVariant::Leave 738 | } 739 | 740 | pub fn bb() -> NodeVariant::BrEntry { 741 | NodeVariant::BrEntry 742 | } 743 | 744 | pub fn bcall() -> NodeVariant::BrCall { 745 | NodeVariant::BrCall 746 | } 747 | 748 | pub fn bif() -> NodeVariant::BrIf { 749 | NodeVariant::BrIf 750 | } 751 | 752 | } 753 | 754 | 755 | -------------------------------------------------------------------------------- /tangle/src/oper.rs: -------------------------------------------------------------------------------- 1 | const HOST_WIDTH: u8 = std::mem::size_of::() as u8; 2 | const STACK_TOP: usize = 0xFFFF_FFFF_FFFF_FFF0; 3 | const STACK: RegSpec = RegSpec::rsp(); 4 | 5 | pub fn check_cond(&mut self, cond: IntCC) -> JitValue { 6 | let flags = match cond { 7 | IntCC::UnsignedGreaterThan => { 8 | vec![Flags::CF, Flags::ZF] 9 | }, 10 | IntCC::UnsignedGreaterThanOrEqual => { 11 | vec![Flags::CF] 12 | }, 13 | IntCC::UnsignedLessThan => { 14 | vec![Flags::CF] 15 | }, 16 | IntCC::UnsignedLessThanOrEqual => { 17 | vec![Flags::CF, Flags::ZF] 18 | }, 19 | IntCC::Equal => { 20 | vec![Flags::ZF] 21 | }, 22 | IntCC::NotEqual => { 23 | vec![Flags::ZF] 24 | }, 25 | _ => unimplemented!("unimplemented check_cond for {}", cond), 26 | }; 27 | let mut val = None; 28 | let mut same = true; 29 | for flag in flags { 30 | let flag_cond = &self.context.flags[flag as usize]; 31 | if let Some(set_val) = val { 32 | if set_val != flag_cond { 33 | same = false; 34 | break; 35 | } 36 | } else { 37 | val = Some(flag_cond); 38 | } 39 | } 40 | if !same { 41 | unimplemented!(); 42 | } 43 | if let Some(JitFlag::Unknown(left, right, res)) = val { 44 | JitValue::Flag(self.builder.ins().ifcmp(*left, *right)) 45 | } else if let Some(JitFlag::Known(c)) = val { 46 | println!("constant eflags {:?} with cond {}", c, cond); 47 | let tmp = match cond { 48 | IntCC::UnsignedGreaterThan => { 49 | // CF=0 and ZF=0 50 | let cf = self.context.check_flag(Flags::CF, false, self.builder); 51 | let zf = self.context.check_flag(Flags::ZF, false, self.builder); 52 | self.band(cf, zf) 53 | }, 54 | IntCC::UnsignedGreaterThanOrEqual => { 55 | // CF=0 56 | self.context.check_flag(Flags::CF, false, self.builder) 57 | }, 58 | IntCC::UnsignedLessThan => { 59 | // CF=1 60 | self.context.check_flag(Flags::CF, true, self.builder) 61 | }, 62 | IntCC::UnsignedLessThanOrEqual => { 63 | // CF=1 or ZF=1 64 | let cf = self.context.check_flag(Flags::CF, true, self.builder); 65 | let zf = self.context.check_flag(Flags::ZF, true, self.builder); 66 | self.bor(cf, zf) 67 | }, 68 | IntCC::Equal => { 69 | self.context.check_flag(Flags::ZF, true, self.builder) 70 | }, 71 | IntCC::NotEqual => { 72 | self.context.check_flag(Flags::ZF, false, self.builder) 73 | } 74 | _ => unimplemented!(), 75 | }; 76 | // XXX: does cranelift do something dumb and actually make this clobber 77 | // an entire register or something? 78 | tmp.zero_extend(self.builder, HOST_WIDTH).into_native(self.builder).unwrap() 79 | } else { 80 | unimplemented!() 81 | } 82 | } 83 | 84 | 85 | /// Get left+right for JitValues. 86 | /// if FLAG is true, then self.context.flags are updated with eflags 87 | pub fn add(&mut self, left: JitTemp, right: JitTemp) -> JitTemp { 88 | use JitTemp::*; 89 | use Flags::*; 90 | let flags = vec![OF, SF, ZF, AF, PF, CF]; 91 | match (left, right) { 92 | (Value(val_left, typ_left), Value(val_right, _)) => { 93 | let (res, flag) = self.builder.ins().iadd_cout(val_left, val_right); 94 | self.context.set_flags(flags, (val_left, val_right, flag)); 95 | Value(res, typ_left) 96 | }, 97 | (left @ _ , right @ Value(_, _)) 98 | | (left @ Value(_, _), right @ _) => { 99 | let typ = Type::int((left.width() * 8) as u16).unwrap(); 100 | let _left = left.into_ssa(typ, self.builder); 101 | let _right = right.zero_extend(self.builder, left.width()).into_ssa(typ, self.builder); 102 | let (res, flag) = self.builder.ins().iadd_cout(_left, _right); 103 | self.context.set_flags(flags, (_left, _right, res)); 104 | Value(res, typ) 105 | }, 106 | (Ref(base, offset), c) if c.is_const() => { 107 | let c = c.into_usize(self.builder).unwrap(); 108 | Ref(base, do_op!(self.context, FLAG, offset, (c), "add {}, {}", flags, +)) 109 | }, 110 | (offset, Ref(base, c)) if offset.is_const() => { 111 | let offset = offset.into_usize(self.builder).unwrap(); 112 | Ref(base, do_op!(self.context, FLAG, offset, (c), "add {}, {}", flags, +)) 113 | }, 114 | (left_c, mut right_c) if left_c.clone().is_const() && right_c.clone().is_const() => { 115 | right_c = right_c.zero_extend(self.builder, left_c.clone().width()); 116 | const_ops!(self.context, FLAG, left_c.clone(), (right_c.clone()), "add {}, {}", flags, +) 117 | }, 118 | (left, right) => unimplemented!("unimplemented add {:?} {:?}", left, right) 119 | } 120 | } 121 | 122 | /// Add with carry 123 | pub fn adc(&mut self, left: JitTemp, right: JitTemp, carry: JitTemp) -> JitTemp { 124 | use JitTemp::*; 125 | use Flags::*; 126 | let flags = vec![OF, SF, ZF, AF, PF, CF]; 127 | match (left.clone(), right.clone(), carry.clone()) { 128 | (Value(val_left, typ_left), Value(val_right, _), Flag(val_carry)) => { 129 | let (res, flag) = self.builder.ins().iadd_carry(val_left, val_right, val_carry); 130 | self.context.set_flags(flags, (val_left, val_right, flag)); 131 | Value(res, typ_left) 132 | }, 133 | (_, _, carry) if carry.is_const() => { 134 | // basically, this sucks a lot: 135 | // iadd_carry is broken on cranelift (#2860), so we work around 136 | // it by lowering to a normal iadd_cout in the CF=false case, 137 | // and to two iadd_cout's in the CF=true case (and then have to 138 | // fixup our carry flag so the spliced +1 doesn't break). 139 | let carry = carry.into_usize(self.builder).unwrap(); 140 | if carry == 0 { 141 | return self.add::(left, right); 142 | } else { 143 | let left_plus_one = self.add::(left, Const8(1)); 144 | // track if the +1 overflowed 145 | let carried = self.context.check_flag(Flags::CF, true, self.builder); 146 | 147 | let res = self.add::(left_plus_one, right); 148 | if(FLAG) { 149 | // if the +1 definitely carried, we know that carry has to 150 | // be set. 151 | println!("carried = {:?}", carried); 152 | if let JitTemp::Const64(1) = carried { 153 | self.context.flags[Flags::CF as usize] = JitFlag::Known(1 << Flags::CF as usize); 154 | } else if let JitTemp::Const64(0) = carried { 155 | (); 156 | } else if let JitTemp::Value(val, typ) = carried { 157 | unimplemented!("double check this"); 158 | // but if we don't know, then we have to OR it with 159 | // the other carry flag (which sucks!) 160 | let other_carried = self.context.check_flag(Flags::CF, true, self.builder); 161 | let maybe_carried = self.bor(other_carried, carried); 162 | self.context.set_flags(vec![Flags::CF], 163 | (left_plus_one.into_ssa(typ, self.builder), 164 | right.into_ssa(typ, self.builder), 165 | maybe_carried.into_ssa(typ, self.builder))); 166 | } else { 167 | unimplemented!(); 168 | } 169 | } 170 | return res; 171 | } 172 | }, 173 | (left @ Value(_, _), right @ _, carry @ _) | 174 | (left @ _, right @ Value(_, _), carry @ _) | 175 | (left @ _, right @ _, carry @ Flag(_)) => { 176 | let typ = Type::int((left.width() * 8) as u16).unwrap(); 177 | let _left = left.into_ssa(typ, self.builder); 178 | let _right = right.zero_extend(self.builder, left.width()).into_ssa(typ, self.builder); 179 | let mut _carry = carry.into_ssa(types::IFLAGS, self.builder); 180 | let (res, flag) = self.builder.ins().iadd_ifcarry(_left, _right, _carry); 181 | self.context.set_flags(flags, (_left, _right, flag)); 182 | Value(res, typ) 183 | }, 184 | (Ref(base, offset), off, carry) if off.is_const() && carry.is_const() => { 185 | let off = off.into_usize(self.builder).unwrap(); 186 | let carry = carry.into_usize(self.builder).unwrap(); 187 | assert_eq!(HOST_WIDTH, 8); 188 | if carry == 0 { 189 | Ref(base, do_op!(self.context, FLAG, offset, off, "add {}, {}", flags, +)) 190 | } else { 191 | Ref(base, do_op!(self.context, FLAG, offset, off, "stc; adc {}, {}", flags, +)) 192 | } 193 | }, 194 | (offset, Ref(base, off), carry) if offset.is_const() && carry.is_const() => { 195 | let offset = offset.into_usize(self.builder).unwrap(); 196 | let carry = carry.into_usize(self.builder).unwrap(); 197 | assert_eq!(HOST_WIDTH, 8); 198 | if carry == 0 { 199 | Ref(base, do_op!(self.context, FLAG, offset, off, "add {}, {}", flags, +)) 200 | } else { 201 | Ref(base, do_op!(self.context, FLAG, offset, off, "stc; adc {}, {}", flags, +)) 202 | } 203 | }, 204 | (left_c, mut right_c, mut carry_c) 205 | if left_c.clone().is_const() 206 | && right_c.clone().is_const() 207 | && carry_c.clone().is_const() => { 208 | right_c = right_c.zero_extend(self.builder, left_c.clone().width()); 209 | let carry_c = carry_c.into_usize(self.builder).unwrap(); 210 | assert_eq!(HOST_WIDTH, 8); 211 | if carry_c == 0 { 212 | const_ops!(self.context, FLAG, left_c.clone(), right_c.clone(), "add {}, {}", flags, +) 213 | } else { 214 | const_ops!(self.context, FLAG, left_c.clone(), right_c.clone(), "stc; adc {}, {}", flags, +) 215 | } 216 | }, 217 | (left, right, carry) => unimplemented!("unimplemented adc {:?} {:?} {:?}", left, right, carry) 218 | } 219 | } 220 | 221 | 222 | 223 | pub fn sub(&mut self, left: JitTemp, right: JitTemp) -> JitTemp { 224 | use JitTemp::*; 225 | use Flags::*; 226 | let flags = vec![OF, SF, ZF, AF, PF, CF]; 227 | match (left, right) { 228 | // sub dest, val is always a dest sized output 229 | (Value(val_left, typ_left), Value(val_right, _)) => { 230 | let (res, flag) = self.builder.ins().isub_ifbout(val_left, val_right); 231 | self.context.set_flags(flags, (val_left, val_right, flag)); 232 | Value(res, typ_left) 233 | }, 234 | (left @ _ , right @ Value(_, _)) 235 | | (left @ Value(_, _), right @ _) => { 236 | let typ = Type::int((left.width() * 8) as u16).unwrap(); 237 | let _left = left.into_ssa(typ, self.builder); 238 | let _right = right.zero_extend(self.builder, left.width()).into_ssa(typ, self.builder); 239 | let (res, flag) = self.builder.ins().isub_ifbout(_left, _right); 240 | self.context.set_flags(flags, (_left, _right, flag)); 241 | Value(res, typ) 242 | }, 243 | // XXX: this is technically incorrect: flags will be wrong, because 244 | // it's actually supposed to be the flags for (base+offset) - c. 245 | // let's just miscompile for now? not sure how common cmp reg, c 246 | // to test pointer equality would be. 247 | (Ref(base, offset), c) if c.is_const() => { 248 | let c = c.into_usize(self.builder).unwrap(); 249 | assert!(c < offset && c < 0x100); 250 | Ref(base, do_op!(self.context, FLAG, offset, c, "sub {}, {}", flags, -)) 251 | }, 252 | (offset, Ref(base, c)) if offset.is_const() => { 253 | let offset = offset.into_usize(self.builder).unwrap(); 254 | assert!(c < offset && c < 0x100); 255 | Ref(base, do_op!(self.context, FLAG, offset, c, "sub {}, {}", flags, -)) 256 | }, 257 | (left_c, mut right_c) if left_c.clone().is_const() && right_c.clone().is_const() => { 258 | right_c = right_c.zero_extend(self.builder, left_c.clone().width()); 259 | const_ops!(self.context, FLAG, left_c.clone(), right_c.clone(), "sub {}, {}", flags, -) 260 | }, 261 | (left, right) => unimplemented!("unimplemented sub {:?} {:?}", left, right) 262 | } 263 | } 264 | 265 | /// Sub with borrow 266 | pub fn sbb(&mut self, left: JitTemp, right: JitTemp, borrow: JitTemp) -> JitTemp { 267 | use JitTemp::*; 268 | use Flags::*; 269 | let flags = vec![OF, SF, ZF, AF, PF, CF]; 270 | match (left.clone(), right.clone(), borrow.clone()) { 271 | (Value(val_left, typ_left), Value(val_right, _), Flag(val_borrow)) => { 272 | let (res, flag) = self.builder.ins().isub_borrow(val_left, val_right, val_borrow); 273 | self.context.set_flags(flags, (val_left, val_right, flag)); 274 | Value(res, typ_left) 275 | }, 276 | (_, _, borrow) if borrow.is_const() => { 277 | // see note in self.adc 278 | let borrow = borrow.into_usize(self.builder).unwrap(); 279 | if borrow == 0 { 280 | return self.sub::(left, right); 281 | } else { 282 | unimplemented!() 283 | } 284 | }, 285 | (left @ Value(_, _), right @ _, borrow @ _) | 286 | (left @ _, right @ Value(_, _), borrow @ _) | 287 | (left @ _, right @ _, borrow @ Flag(_)) => { 288 | let typ = Type::int((left.width() * 8) as u16).unwrap(); 289 | let _left = left.into_ssa(typ, self.builder); 290 | let _right = right.zero_extend(self.builder, left.width()).into_ssa(typ, self.builder); 291 | let mut _borrow = borrow.into_ssa(types::IFLAGS, self.builder); 292 | let (res, flag) = self.builder.ins().isub_borrow(_left, _right, _borrow); 293 | self.context.set_flags(flags, (_left, _right, flag)); 294 | Value(res, typ) 295 | }, 296 | (Ref(base, offset), off, borrow) if off.is_const() && borrow.is_const() => { 297 | let off = off.into_usize(self.builder).unwrap(); 298 | let borrow = borrow.into_usize(self.builder).unwrap(); 299 | assert_eq!(HOST_WIDTH, 8); 300 | if borrow == 0 { 301 | Ref(base, do_op!(self.context, FLAG, offset, off, "sub {}, {}", flags, +)) 302 | } else { 303 | Ref(base, do_op!(self.context, FLAG, offset, off, "stc; sbb {}, {}", flags, +)) 304 | } 305 | }, 306 | (offset, Ref(base, off), borrow) if offset.is_const() && borrow.is_const() => { 307 | let offset = offset.into_usize(self.builder).unwrap(); 308 | let borrow = borrow.into_usize(self.builder).unwrap(); 309 | assert_eq!(HOST_WIDTH, 8); 310 | if borrow == 0 { 311 | Ref(base, do_op!(self.context, FLAG, offset, off, "sub {}, {}", flags, +)) 312 | } else { 313 | Ref(base, do_op!(self.context, FLAG, offset, off, "stc; sbb {}, {}", flags, +)) 314 | } 315 | }, 316 | (left_c, mut right_c, mut borrow) 317 | if left_c.clone().is_const() 318 | && right_c.clone().is_const() 319 | && borrow.clone().is_const() => { 320 | right_c = right_c.zero_extend(self.builder, left_c.clone().width()); 321 | let borrow = borrow.into_usize(self.builder).unwrap(); 322 | assert_eq!(HOST_WIDTH, 8); 323 | if borrow == 0 { 324 | const_ops!(self.context, FLAG, left_c.clone(), right_c.clone(), "sub {}, {}", flags, +) 325 | } else { 326 | const_ops!(self.context, FLAG, left_c.clone(), right_c.clone(), "stc; sbb {}, {}", flags, +) 327 | } 328 | }, 329 | (left, right, borrow) => unimplemented!("unimplemented sbb {:?} {:?} {:?}", left, right, borrow) 330 | } 331 | } 332 | 333 | 334 | 335 | pub fn band(&mut self, mut left: JitTemp, mut right: JitTemp) -> JitTemp { 336 | use JitTemp::*; 337 | let biggest = max(left.width(), right.width()); 338 | left = left.zero_extend(self.builder, biggest); 339 | right = right.zero_extend(self.builder, biggest); 340 | // XXX: set flags!! 341 | match (left.clone(), right.clone()) { 342 | (c, other) | (other, c) if c.is_const() => { 343 | let fixed_c = c.into_usize(self.builder).unwrap(); 344 | match (fixed_c, other.clone()) { 345 | //0 & val or val & 0 = 0 346 | (0, _) => return JitTemp::val_of_width(0, biggest), 347 | // !0 & val or val & !0 = val 348 | (const { !0 as usize }, _) => return other, 349 | // left_c & right_c is just left & right 350 | (fixed_c, other) if other.clone().is_const() => { 351 | let fixed_other = other.into_usize(self.builder).unwrap(); 352 | return JitTemp::val_of_width(fixed_c & fixed_other, biggest) 353 | }, 354 | (_, _) => (), 355 | } 356 | }, 357 | _ => (), 358 | }; 359 | 360 | match (left, right) { 361 | // add rsp, 0x8 becomes a redefine of rsp with offset -= 1 362 | (Value(val_left, typ_left), Value(val_right, typ_right)) => { 363 | Value(self.builder.ins().band(val_left, val_right), typ_left) 364 | }, 365 | (val @ _, Value(ssa, ssa_typ)) 366 | | (Value(ssa, ssa_typ), val @ _) => { 367 | let ssa_val = val.into_ssa(self.int, self.builder); 368 | Value(self.builder.ins().band(ssa_val, ssa), ssa_typ) 369 | }, 370 | (Ref(base, offset), _) 371 | | (_, Ref(base, offset)) => { 372 | unimplemented!("band pointer") 373 | }, 374 | (left, right) => unimplemented!("unimplemented band {:?} {:?}", left, right) 375 | } 376 | } 377 | 378 | pub fn bor(&mut self, mut left: JitTemp, mut right: JitTemp) -> JitTemp { 379 | use JitTemp::*; 380 | let biggest = max(left.width(), right.width()); 381 | left = left.zero_extend(self.builder, biggest); 382 | right = right.zero_extend(self.builder, biggest); 383 | // XXX: set flags!! 384 | match (left.clone(), right.clone()) { 385 | (c, other) | (other, c) if c.is_const() => { 386 | let fixed_c = c.into_usize(self.builder).unwrap(); 387 | match (fixed_c, other) { 388 | //0 | val or val | 0 = 0 389 | (0, val) => return val, 390 | // left_c | right_c is just left | right 391 | (fixed_c, other) if other.is_const() => { 392 | let fixed_other = other.into_usize(self.builder).unwrap(); 393 | return JitTemp::val_of_width(fixed_c | fixed_other, biggest) 394 | }, 395 | (_, _) => (), 396 | } 397 | }, 398 | _ => (), 399 | }; 400 | 401 | match (left, right) { 402 | (Value(val_left, typ_left), Value(val_right, typ_right)) => { 403 | Value(self.builder.ins().bor(val_left, val_right), typ_left) 404 | }, 405 | (val @ _, Value(ssa, ssa_typ)) 406 | | (Value(ssa, ssa_typ), val @ _) => { 407 | let ssa_val = val.into_ssa(self.int, self.builder); 408 | Value(self.builder.ins().bor(ssa_val, ssa), ssa_typ) 409 | }, 410 | (Ref(base, offset), _) 411 | | (_, Ref(base, offset)) => { 412 | unimplemented!("bor pointer") 413 | }, 414 | (left, right) => unimplemented!("unimplemented bor {:?} {:?}", left, right) 415 | } 416 | } 417 | 418 | pub fn bxor(&mut self, mut left: JitTemp, mut right: JitTemp) -> JitTemp { 419 | use JitTemp::*; 420 | if left == right { 421 | return JitTemp::val_of_width(0, left.width()) 422 | } 423 | let biggest = max(left.width(), right.width()); 424 | left = left.zero_extend(self.builder, biggest); 425 | right = right.zero_extend(self.builder, biggest); 426 | // XXX: set flags 427 | match (left.clone(), right.clone()) { 428 | (c, other) | (other, c) if c.is_const() => { 429 | let fixed_c = c.into_usize(self.builder).unwrap(); 430 | match (fixed_c, other) { 431 | (fixed_c, other) if other.is_const() => { 432 | let fixed_other = other.into_usize(self.builder).unwrap(); 433 | return JitTemp::val_of_width(fixed_c ^ fixed_other, biggest) 434 | }, 435 | (_, _) => (), 436 | } 437 | }, 438 | _ => (), 439 | } 440 | 441 | match (left, right) { 442 | (Value(val_left, typ_left), Value(val_right, typ_right)) => { 443 | Value(self.builder.ins().bxor(val_left, val_right), typ_left) 444 | }, 445 | (val @ _, Value(ssa, ssa_typ)) 446 | | (Value(ssa, ssa_typ), val @ _) => { 447 | let ssa_val = val.into_ssa(self.int, self.builder); 448 | Value(self.builder.ins().bxor(ssa_val, ssa), ssa_typ) 449 | }, 450 | (left, right) => unimplemented!("unimplemented bxor {:?} {:?}", left, right) 451 | } 452 | } 453 | 454 | pub fn bitselect(&mut self, mut control: JitTemp, mut left: JitTemp, mut right: JitTemp) -> JitTemp { 455 | use JitTemp::*; 456 | let biggest = max(control.width(), max(left.width(), right.width())); 457 | println!("bitselect {:?} {:?} {:?} (biggest={})", control, left, right, biggest); 458 | control = control.zero_extend(self.builder, biggest); 459 | left = left.zero_extend(self.builder, biggest); 460 | right = right.zero_extend(self.builder, biggest); 461 | // bitselect doesn't have a set flags version 462 | if control.is_const() { 463 | let fixed_control = control.clone().into_usize(self.builder).unwrap(); 464 | match fixed_control { 465 | //bitselect(0, left, right) => right 466 | 0 => return right, 467 | //bitselect(!0, left, right) => left 468 | const { !0 as usize } => return left, 469 | _ => (), 470 | }; 471 | } 472 | if left.is_const() { 473 | let fixed_left = left.clone().into_usize(self.builder).unwrap(); 474 | match fixed_left { 475 | //bitselect(control, 0, right) => !control & right 476 | 0 => { 477 | let ncontrol = self.bnot(control); 478 | return self.band(ncontrol, right) 479 | }, 480 | _ => () 481 | }; 482 | } 483 | if right.is_const() { 484 | let fixed_right = right.clone().into_usize(self.builder).unwrap(); 485 | match fixed_right { 486 | //bitselect(control, left, 0) => control & left 487 | 0 => return self.band(control, left), 488 | _ => () 489 | } 490 | } 491 | match (control, left, right) { 492 | (control, left, right) if control.is_const() && left.is_const() && right.is_const() => { 493 | let true_mask = self.band(control.clone(), left); 494 | let not_control = self.bnot(control); 495 | let false_mask = self.band(not_control, right); 496 | self.bor(true_mask, false_mask) 497 | }, 498 | (control @ _, Value(val_left, typ_left), Value(val_right, _)) => { 499 | let control = control.into_ssa(self.int, self.builder); 500 | Value(self.builder.ins().bitselect(control, val_left, val_right), typ_left) 501 | }, 502 | (control, left, right) => 503 | unimplemented!("unimplemented bitselect {:?} {:?} {:?}", control, left, right) 504 | 505 | } 506 | } 507 | 508 | pub fn bnot(&mut self, val: JitTemp) -> JitTemp { 509 | use JitTemp::*; 510 | match val { 511 | Value(val, typ) => Value(self.builder.ins().bnot(val), typ), 512 | Const8(c) => Const8(!c), 513 | Const16(c) => Const16(!c), 514 | Const32(c) => Const32(!c), 515 | Const64(c) => Const64(!c), 516 | _ => unimplemented!("bnot {:?}", val), 517 | } 518 | } 519 | 520 | -------------------------------------------------------------------------------- /tangle/src/opt.rs: -------------------------------------------------------------------------------- 1 | use crate::node::{Node, Operation, NodeVariant}; 2 | use crate::port::{Port, PortMeta, Edge, EdgeVariant}; 3 | use crate::region::{Region, NodeGraph, PortGraph}; 4 | 5 | use crate::petgraph::visit::{IntoEdgeReferences, EdgeRef}; 6 | 7 | //trait Pass where 8 | // F: Pass 9 | //{ 10 | // type First: Pass = F; 11 | // type Next: Pass; 12 | // fn visit_node(nodes: &NodeGraph, node: Node) -> PassEffect; 13 | //} 14 | // 15 | //enum PassEffect { 16 | // Return(Node), 17 | //} 18 | // 19 | //struct ConstantAdd; 20 | //impl> Pass for ConstantAdd { 21 | // // {Constant, Constant} -> Add -> X; Constant -> X 22 | // // incoming count for input ports is always 1 23 | // // really we want a dependantly typed Node.become function that forces you 24 | // // to provide a mapping function from Variant ports->Variant ports 25 | // type Next = Nop; 26 | // fn visit_node(nodes: &NodeGraph, node: Node) -> PassEffect { 27 | // // do something 28 | // let mut add_node = node.as_variant::(); 29 | // if let Some(add_node) = add_node && add_node.0 == Operation::Add { 30 | // println!("add"); 31 | // } 32 | // 33 | // if 1 == 1 { 34 | // return >::visit_node(nodes, node) 35 | // } else { 36 | // return >::visit_node(nodes, node) 37 | // } 38 | // } 39 | //} 40 | // 41 | //struct Nop; 42 | //impl> Pass for Nop { 43 | // type Next = (); 44 | // fn visit_node(nodes: &NodeGraph, node: Node) -> PassEffect { 45 | // return PassEffect::Return(node); 46 | // } 47 | //} 48 | // 49 | //impl> Pass for () { 50 | // type Next = (); 51 | // fn visit_node(nodes: &NodeGraph, node: Node) -> PassEffect { 52 | // panic!(); 53 | // } 54 | //} 55 | // 56 | //struct Fusor; 57 | //impl Pass for Fusor { 58 | // type First = Fusor; 59 | // type Next = ConstantAdd; 60 | // fn visit_node(nodes: &NodeGraph, node: Node) -> PassEffect { 61 | // return >::visit_node(nodes, node) 62 | // } 63 | //} 64 | // 65 | //pub struct PassRunner; 66 | //impl PassRunner { 67 | // pub fn run(&self, region: &mut Region) { 68 | // let Region { nodes, ports, .. } = region; 69 | // 70 | //let mut nodes_weights = nodes.node_indices().collect::>(); 71 | //while let Some(node) = nodes_weights.pop() { 72 | // let mut tmp_node = Node::new(box Node::simple(Operation::Nop)); 73 | // core::mem::swap(&mut tmp_node, &mut nodes[node]); 74 | // let res = Fusor::visit_node(nodes, tmp_node); 75 | // match res { 76 | // PassEffect::Return(mut n) => 77 | // core::mem::swap(&mut n, &mut nodes[node]), 78 | // _ => panic!(), 79 | // } 80 | //} 81 | // } 82 | //} 83 | -------------------------------------------------------------------------------- /tangle/src/port.rs: -------------------------------------------------------------------------------- 1 | use petgraph::graph::{NodeIndex, EdgeIndex}; 2 | use frunk::prelude::*; 3 | use frunk::hlist::*; 4 | use frunk::*; 5 | use yaxpeax_x86::long_mode::RegSpec; 6 | 7 | use core::ops::{Deref}; 8 | 9 | use crate::node::{NodeIdx}; 10 | use crate::time::Timestamp; 11 | 12 | #[derive(Hash, PartialOrd, Ord, PartialEq, Eq, Clone, Copy, Default)] 13 | pub struct PortIdxToken(u32); 14 | pub type PortIdx = NodeIndex; 15 | unsafe impl petgraph::matrix_graph::IndexType for PortIdxToken { 16 | #[inline(always)] 17 | fn new(x: usize) -> Self { 18 | PortIdxToken(x as u32) 19 | } 20 | #[inline(always)] 21 | fn index(&self) -> usize { 22 | self.0 as usize 23 | } 24 | #[inline(always)] 25 | fn max() -> Self { 26 | PortIdxToken(::std::u32::MAX) 27 | } 28 | } 29 | 30 | impl core::fmt::Debug for PortIdxToken { 31 | fn fmt(&self, fmt: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> { 32 | write!(fmt, "p{}", self.0) 33 | } 34 | } 35 | 36 | #[derive(PartialEq, Debug, Copy, Clone)] 37 | pub enum EdgeVariant { 38 | State, 39 | Data, 40 | } 41 | 42 | #[derive(Clone, Copy, PartialEq, Debug, Eq, Hash)] 43 | pub enum Storage { 44 | Virtual(u16), // todo: width? 45 | Physical(RegSpec), 46 | Immaterial(Option), // used for constant operands and state edges - does not receive a vreg 47 | } 48 | 49 | impl core::fmt::Display for Storage { 50 | fn fmt(&self, fmt: &mut core::fmt::Formatter) -> Result<(), std::fmt::Error> { 51 | match self { 52 | Storage::Virtual(v) => write!(fmt, "v{}", v), 53 | Storage::Physical(p) => write!(fmt, "{}", p), 54 | Storage::Immaterial(_) => write!(fmt, "imm"), 55 | } 56 | } 57 | } 58 | 59 | #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] 60 | pub struct OptionalStorage(pub Option); 61 | impl core::fmt::Display for OptionalStorage { 62 | fn fmt(&self, fmt: &mut core::fmt::Formatter) -> Result<(), std::fmt::Error> { 63 | match self.0 { 64 | None => write!(fmt, "()"), 65 | Some(r) => write!(fmt, "{}", r), 66 | } 67 | } 68 | } 69 | 70 | impl Deref for OptionalStorage { 71 | type Target = Option; 72 | fn deref(&self) -> &::Target { 73 | &self.0 74 | } 75 | } 76 | 77 | 78 | 79 | pub mod PortMeta { 80 | use frunk::HList; 81 | #[derive(Debug, Clone)] 82 | pub struct Constant(pub isize); 83 | #[derive(Debug, Clone)] 84 | pub struct StackOff(pub i32); 85 | pub type All = HList![Option, Option]; 86 | } 87 | 88 | #[derive(Clone, Debug)] 89 | pub struct Port { 90 | pub id: u32, 91 | pub variant: EdgeVariant, 92 | pub meta: Option, 93 | pub storage: OptionalStorage, 94 | pub time: Option, 95 | pub done: bool, 96 | // The node the port is attached to. If none, it isn't attached to any (e.g. a region 97 | // source/sink) 98 | pub node: Option, 99 | } 100 | impl Port { 101 | pub fn new() -> Self { 102 | let mut guard = PORT_COUNT.lock().unwrap(); 103 | let curr = *guard; 104 | *guard += 1; 105 | println!("created port v{:?}", curr); 106 | Port { 107 | id: curr, 108 | variant: EdgeVariant::Data, 109 | storage: OptionalStorage(None), 110 | time: None, 111 | done: false, 112 | node: None, 113 | meta: None, 114 | } 115 | } 116 | 117 | pub fn set_variant(&mut self, var: EdgeVariant) { 118 | self.variant = var; 119 | } 120 | 121 | pub fn set_storage(&mut self, store: Storage) { 122 | //assert!(self.variant != EdgeVariant::State, "tried to set storage on a state port"); 123 | if let Some(exist) = self.storage.0 { 124 | if exist != store { 125 | println!("port already has storage {}, setting to {}", exist, store); 126 | } 127 | } 128 | self.storage = OptionalStorage(Some(store)); 129 | } 130 | 131 | pub fn set_meta<'this, 'a, K, T>(&'this mut self, meta_val: K) where 132 | PortMeta::All: LiftFrom, T>, 133 | >::Output: Plucker<&'a mut Option, T>, 134 | Option: core::fmt::Debug, 135 | 'this: 'a, 136 | K: 'static, 137 | { 138 | if let None = self.meta { 139 | let x: Option = Some(Some(meta_val).lift_into()); 140 | self.meta = x; 141 | } else { 142 | self.meta.as_mut().map(|mut m| { 143 | let (head, _): (&'a mut Option, _) = m.to_mut().pluck(); 144 | *head = Some(meta_val); 145 | }); 146 | } 147 | } 148 | 149 | pub fn get_meta<'this, 'a, K, T>(&'this self) -> Option<&'a K> where 150 | >::Output: Plucker<&'this Option, T>, 151 | K: 'static + Clone, 152 | 'this: 'a, 153 | { 154 | self.meta.as_ref().and_then(|m| { 155 | let (head, _) = m.to_ref().pluck::<&'this Option, _>(); 156 | head.as_ref() 157 | }) 158 | } 159 | 160 | pub fn set_node(&mut self, n: NodeIdx) { 161 | self.node = Some(n); 162 | } 163 | } 164 | 165 | use std::sync::Mutex; 166 | lazy_static::lazy_static! { 167 | static ref PORT_COUNT: Mutex = Mutex::new(0); 168 | } 169 | 170 | #[derive(Clone, Debug)] 171 | pub struct Edge { 172 | pub variant: EdgeVariant, 173 | } 174 | pub type PortEdge = EdgeIndex; 175 | 176 | 177 | -------------------------------------------------------------------------------- /tangle/src/time.rs: -------------------------------------------------------------------------------- 1 | // We need a timestamp for nodes and ports such that we can push and pull the in time 2 | // relative to other instructions in the same timeslice, but without invalidating e.g. 3 | // the dependency ordering that we primarily use the times for - that is, an instruction 4 | // at t=2 is always emitted after its t=1 dependencies. 5 | // 6 | // We use a fixednum timestamp for this. The 'major' time is the dependency time, 7 | // while the 'minor' time allows us to order instructions within it; overflowing 8 | // the minor time will just panic for now. 9 | 10 | 11 | use rangemap::StepLite; 12 | 13 | #[derive(Clone, Copy, Debug, PartialEq, Eq, Ord)] 14 | pub struct Timestamp { 15 | pub major: u16, 16 | pub minor: u16, 17 | } 18 | 19 | impl Timestamp { 20 | pub fn new() -> Self { 21 | Timestamp { 22 | major: 0, 23 | minor: 8, 24 | } 25 | } 26 | 27 | /// Increment the timestamp's major time 28 | #[must_use] 29 | pub fn increment(mut self) -> Self { 30 | self.major = self.major.checked_add(1).unwrap(); 31 | self 32 | } 33 | 34 | /// Decrement the timestamp's major time 35 | #[must_use] 36 | pub fn decrement(mut self) -> Self { 37 | self.major = self.major.checked_sub(1).unwrap(); 38 | self 39 | } 40 | 41 | /// Push the timestamp to a later minor time 42 | #[must_use] 43 | pub fn push(mut self) -> Self { 44 | self.minor = self.minor.checked_add(1).unwrap(); 45 | self 46 | } 47 | 48 | /// Pull the timestamp to a sooner minor time 49 | #[must_use] 50 | pub fn pull(mut self) -> Self { 51 | self.minor = self.minor.checked_sub(1).unwrap(); 52 | self 53 | } 54 | } 55 | 56 | impl PartialOrd for Timestamp { 57 | fn partial_cmp(&self, other: &Timestamp) -> Option { 58 | Some(self.major.cmp(&other.major).then(self.minor.cmp(&other.minor))) 59 | } 60 | } 61 | 62 | #[cfg(test)] 63 | mod test { 64 | use super::*; 65 | 66 | #[test] 67 | fn timestamp_cmp() { 68 | assert!(Timestamp::new() == Timestamp::new()); 69 | assert!(Timestamp::new() <= Timestamp::new()); 70 | assert!(Timestamp::new() >= Timestamp::new()); 71 | assert!(Timestamp::new().increment() > Timestamp::new()); 72 | assert!(Timestamp::new().push() > Timestamp::new()); 73 | assert!(Timestamp::new().push() < Timestamp::new().increment()); 74 | assert!(Timestamp::new().increment().push() > Timestamp::new().push()); 75 | assert!(Timestamp::new().increment().push() > Timestamp::new().increment()); 76 | } 77 | } 78 | 79 | impl StepLite for Timestamp { 80 | fn add_one(&self) -> Self { 81 | if self.minor == u16::MAX { let mut new = self.increment(); new.minor = u16::MIN; new } 82 | else { self.push() } 83 | } 84 | 85 | fn sub_one(&self) -> Self { 86 | if self.minor == u16::MIN { let mut new = self.decrement(); new.minor = u16::MAX; new } 87 | else { self.pull() } 88 | } 89 | } 90 | 91 | use core::fmt::{Display, Formatter}; 92 | impl Display for Timestamp { 93 | fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), std::fmt::Error> { 94 | write!(fmt, "{}.{}", self.major, self.minor) 95 | } 96 | } 97 | 98 | /// Vector clock used for state effect ordering. 99 | struct Clock { 100 | members: Vec 101 | } 102 | 103 | impl Clock { 104 | pub fn new() -> Self { 105 | Clock { members: vec![] } 106 | } 107 | 108 | #[must_use] 109 | pub fn increment(mut self, member: usize) -> Self { 110 | self.members[member] = self.members[member].increment(); 111 | self 112 | } 113 | 114 | #[must_use] 115 | pub fn decrement(mut self, member: usize) -> Self { 116 | self.members[member] = self.members[member].decrement(); 117 | self 118 | } 119 | 120 | } 121 | 122 | 123 | 124 | 125 | 126 | 127 | 128 | 129 | 130 | 131 | 132 | 133 | 134 | 135 | 136 | 137 | 138 | --------------------------------------------------------------------------------