nac3/nac3standalone/src/main.rs

331 lines
10 KiB
Rust

use std::time::SystemTime;
use std::{collections::HashSet, fs};
use inkwell::{
passes::{PassManager, PassManagerBuilder},
targets::*,
OptimizationLevel,
};
use parking_lot::RwLock;
use rustpython_parser::{
ast::{fold::Fold, StmtKind},
parser,
};
use std::{cell::RefCell, collections::HashMap, path::Path, sync::Arc};
use nac3core::{
codegen::{CodeGenTask, WithCall, WorkerRegistry},
symbol_resolver::SymbolResolver,
toplevel::{DefinitionId, FunInstance, TopLevelComposer, TopLevelContext, TopLevelDef},
typecheck::{
type_inferencer::{FunctionData, Inferencer},
typedef::{FunSignature, FuncArg, TypeEnum},
},
};
mod basic_symbol_resolver;
fn main() {
Target::initialize_all(&InitializationConfig::default());
let program = match fs::read_to_string("mandelbrot.py") {
Ok(program) => program,
Err(err) => {
println!("Cannot open input file: {}", err);
return;
}
};
let start = SystemTime::now();
let composer = TopLevelComposer::new();
let mut unifier = composer.unifier.clone();
let primitives = composer.primitives_ty;
let top_level = Arc::new(composer.make_top_level_context());
unifier.top_level = Some(top_level.clone());
let output_fun = unifier.add_ty(TypeEnum::TFunc(RefCell::new(FunSignature {
args: vec![FuncArg {
name: "c".into(),
ty: primitives.int32,
default_value: None,
}],
ret: primitives.none,
vars: HashMap::new(),
})));
let output_id = top_level.definitions.read().len();
top_level
.definitions
.write()
.push(Arc::new(RwLock::new(TopLevelDef::Function {
name: "output".into(),
signature: output_fun,
instance_to_stmt: HashMap::new(),
instance_to_symbol: [("".to_string(), "output".to_string())]
.iter()
.cloned()
.collect(),
var_id: Default::default(),
resolver: None,
})));
// dummy resolver...
let resolver = Arc::new(Box::new(basic_symbol_resolver::Resolver {
id_to_type: HashMap::new(),
id_to_def: HashMap::new(),
class_names: Default::default(),
}) as Box<dyn SymbolResolver + Send + Sync>);
let mut functions = HashMap::new();
for stmt in parser::parse_program(&program).unwrap().into_iter() {
if let StmtKind::FunctionDef {
name,
body,
args,
returns,
..
} = stmt.node
{
let args = args
.args
.into_iter()
.map(|arg| FuncArg {
name: arg.node.arg.to_string(),
ty: resolver
.parse_type_annotation(
&top_level.definitions.read(),
&mut unifier,
&primitives,
&arg.node
.annotation
.expect("expected type annotation in parameters"),
)
.unwrap(),
default_value: None,
})
.collect();
let ret = returns
.map(|r| {
resolver
.parse_type_annotation(
&top_level.definitions.read(),
&mut unifier,
&primitives,
&r,
)
.unwrap()
})
.unwrap_or(primitives.none);
let signature = FunSignature {
args,
ret,
vars: Default::default(),
};
let fun_ty = unifier.add_ty(TypeEnum::TFunc(RefCell::new(signature.clone())));
let id = top_level.definitions.read().len();
top_level
.definitions
.write()
.push(Arc::new(RwLock::new(TopLevelDef::Function {
name: name.clone(),
signature: fun_ty,
var_id: vec![],
instance_to_stmt: HashMap::new(),
instance_to_symbol: HashMap::new(),
resolver: None,
})));
functions.insert(name, (id, body, signature));
} else {
panic!("unsupported statement type");
}
}
let setup_time = SystemTime::now();
println!(
"Setup time: {}ms",
setup_time
.duration_since(start)
.unwrap()
.as_millis()
);
let mut id_to_def: HashMap<_, _> = functions
.iter()
.map(|(k, v)| (k.clone(), DefinitionId(v.0)))
.collect();
id_to_def.insert("output".into(), DefinitionId(output_id));
let mut id_to_type: HashMap<_, _> = functions
.iter()
.map(|(k, v)| {
(
k.clone(),
unifier.add_ty(TypeEnum::TFunc(RefCell::new(v.2.clone()))),
)
})
.collect();
id_to_type.insert("output".into(), output_fun);
let resolver = Arc::new(Box::new(basic_symbol_resolver::Resolver {
class_names: Default::default(),
id_to_type,
id_to_def,
}) as Box<dyn SymbolResolver + Send + Sync>);
for (_, (id, ast, signature)) in functions.into_iter() {
if let TopLevelDef::Function {
resolver: r,
instance_to_stmt,
..
} = &mut *top_level.definitions.read()[id].write()
{
*r = Some(resolver.clone());
let return_type = if unifier.unioned(primitives.none, signature.ret) {
None
} else {
Some(signature.ret)
};
let mut function_data = FunctionData {
resolver: resolver.clone(),
bound_variables: Vec::new(),
return_type,
};
let mut virtual_checks = Vec::new();
let mut calls = HashMap::new();
let mut identifiers = HashSet::new();
let mut variable_mapping = HashMap::new();
for arg in signature.args.iter() {
identifiers.insert(arg.name.clone());
variable_mapping.insert(arg.name.clone(), arg.ty);
}
let mut inferencer = Inferencer {
top_level: &top_level,
function_data: &mut function_data,
unifier: &mut unifier,
variable_mapping,
primitives: &primitives,
virtual_checks: &mut virtual_checks,
calls: &mut calls,
defined_identifiers: identifiers.clone(),
};
let statements = ast
.into_iter()
.map(|v| inferencer.fold_stmt(v))
.collect::<Result<Vec<_>, _>>()
.unwrap();
let returned = inferencer
.check_block(&statements, &mut identifiers)
.unwrap();
if return_type.is_some() && !returned {
panic!("expected return");
}
instance_to_stmt.insert(
"".to_string(),
FunInstance {
body: statements,
unifier_id: 0,
calls,
subst: Default::default(),
},
);
}
}
let inference_time = SystemTime::now();
println!(
"Type inference time: {}ms",
inference_time
.duration_since(setup_time)
.unwrap()
.as_millis()
);
let top_level = Arc::new(TopLevelContext {
definitions: Arc::new(RwLock::new(std::mem::take(
&mut *top_level.definitions.write(),
))),
unifiers: Arc::new(RwLock::new(vec![(
unifier.get_shared_unifier(),
primitives,
)])),
});
let instance = {
let defs = top_level.definitions.read();
let mut instance = defs[resolver.get_identifier_def("run").unwrap().0].write();
if let TopLevelDef::Function {
instance_to_stmt,
instance_to_symbol,
..
} = &mut *instance
{
instance_to_symbol.insert("".to_string(), "run".to_string());
instance_to_stmt[""].clone()
} else {
unreachable!()
}
};
let signature = FunSignature {
args: vec![],
ret: primitives.int32,
vars: HashMap::new(),
};
let task = CodeGenTask {
subst: Default::default(),
symbol_name: "run".to_string(),
body: instance.body,
signature,
resolver,
unifier: top_level.unifiers.read()[instance.unifier_id].clone(),
calls: instance.calls,
};
let f = Arc::new(WithCall::new(Box::new(move |module| {
let codegen_time = SystemTime::now();
println!(
"Code generation time: {}ms",
codegen_time
.duration_since(inference_time)
.unwrap()
.as_millis()
);
let builder = PassManagerBuilder::create();
builder.set_optimization_level(OptimizationLevel::Aggressive);
let passes = PassManager::create(());
builder.populate_module_pass_manager(&passes);
passes.run_on(module);
let triple = TargetMachine::get_default_triple();
let target =
Target::from_triple(&triple).expect("couldn't create target from target triple");
let target_machine = target
.create_target_machine(
&triple,
"",
"",
OptimizationLevel::Default,
RelocMode::Default,
CodeModel::Default,
)
.expect("couldn't create target machine");
target_machine
.write_to_file(module, FileType::Object, Path::new("mandelbrot.o"))
.expect("couldn't write module to file");
println!(
"LLVM time: {}ms",
SystemTime::now()
.duration_since(codegen_time)
.unwrap()
.as_millis()
);
println!("IR:\n{}", module.print_to_string().to_str().unwrap());
})));
let threads = ["test"];
let (registry, handles) = WorkerRegistry::create_workers(&threads, top_level, f);
registry.add_task(task);
registry.wait_tasks_complete(handles);
println!("object file is in mandelbrot.o");
}