2023-09-06 16:34:08 +08:00
use clap ::Parser ;
2021-10-16 22:17:36 +08:00
use inkwell ::{
2022-02-21 18:27:46 +08:00
memory_buffer ::MemoryBuffer ,
2023-09-11 13:14:35 +08:00
passes ::PassBuilderOptions ,
2023-10-13 11:28:30 +08:00
support ::is_multithreaded ,
2021-10-16 22:17:36 +08:00
targets ::* ,
2022-02-21 18:27:46 +08:00
OptimizationLevel ,
2021-10-16 22:17:36 +08:00
} ;
2022-02-21 18:27:46 +08:00
use parking_lot ::{ Mutex , RwLock } ;
2023-11-26 09:09:24 +08:00
use std ::{ collections ::HashMap , fs , path ::Path , sync ::Arc } ;
2023-11-15 17:30:26 +08:00
use std ::collections ::HashSet ;
2020-12-17 22:20:30 +08:00
2021-08-19 15:30:52 +08:00
use nac3core ::{
2021-10-17 13:02:18 +08:00
codegen ::{
2023-09-11 13:12:46 +08:00
concrete_type ::ConcreteTypeStore , irrt ::load_irrt , CodeGenLLVMOptions ,
CodeGenTargetMachineOptions , CodeGenTask , DefaultCodeGenerator , WithCall , WorkerRegistry ,
2021-10-17 13:02:18 +08:00
} ,
2021-08-25 15:30:36 +08:00
symbol_resolver ::SymbolResolver ,
2021-12-01 22:44:53 +08:00
toplevel ::{
2022-02-21 18:27:46 +08:00
composer ::TopLevelComposer , helper ::parse_parameter_default_value , type_annotation ::* ,
TopLevelDef ,
2021-12-01 22:44:53 +08:00
} ,
2022-02-21 18:27:46 +08:00
typecheck ::{
type_inferencer ::PrimitiveStore ,
typedef ::{ FunSignature , Type , Unifier } ,
} ,
} ;
use nac3parser ::{
2023-12-05 14:37:08 +08:00
ast ::{ Constant , Expr , ExprKind , StmtKind , StrRef } ,
2022-02-21 18:27:46 +08:00
parser ,
2021-08-19 15:30:52 +08:00
} ;
2020-12-17 22:20:30 +08:00
2021-08-19 15:30:52 +08:00
mod basic_symbol_resolver ;
2021-09-19 16:19:16 +08:00
use basic_symbol_resolver ::* ;
2023-12-11 14:41:58 +08:00
use nac3core ::toplevel ::composer ::ComposerConfig ;
2020-12-17 22:20:30 +08:00
2023-09-06 16:34:08 +08:00
/// Command-line argument parser definition.
#[ derive(Parser) ]
#[ command(author, version, about, long_about = None) ]
struct CommandLineArgs {
/// The name of the input file.
file_name : String ,
2023-10-13 14:57:09 +08:00
/// The number of threads allocated to processing the source file. If 0 is passed to this
/// parameter, all available threads will be used for compilation.
2023-10-13 11:22:29 +08:00
#[ arg(short = 'T', default_value_t = 1) ]
2023-09-06 16:34:08 +08:00
threads : u32 ,
2023-09-06 18:49:58 +08:00
/// The level to optimize the LLVM IR.
#[ arg(short = 'O', default_value_t = 2, value_parser = clap::value_parser!(u32).range(0..=3)) ]
opt_level : u32 ,
/// Whether to emit LLVM IR at the end of every module.
2023-10-13 15:00:39 +08:00
///
/// If multithreaded compilation is also enabled, each thread will emit its own module.
2023-09-06 18:49:58 +08:00
#[ arg(long, default_value_t = false) ]
emit_llvm : bool ,
2023-09-12 14:41:34 +08:00
/// The target triple to compile for.
#[ arg(long) ]
triple : Option < String > ,
/// The target CPU to compile for.
#[ arg(long) ]
mcpu : Option < String > ,
/// Additional target features to enable/disable, specified using the `+`/`-` prefixes.
#[ arg(long) ]
target_features : Option < String > ,
2023-09-06 16:34:08 +08:00
}
2022-04-18 16:02:48 +08:00
fn handle_typevar_definition (
var : & Expr ,
resolver : & ( dyn SymbolResolver + Send + Sync ) ,
def_list : & [ Arc < RwLock < TopLevelDef > > ] ,
unifier : & mut Unifier ,
primitives : & PrimitiveStore ,
2023-11-15 17:30:26 +08:00
) -> Result < Type , HashSet < String > > {
2023-11-28 17:37:49 +08:00
let ExprKind ::Call { func , args , .. } = & var . node else {
2023-11-15 17:30:26 +08:00
return Err ( HashSet ::from ( [
format! (
" expression {var:?} cannot be handled as a generic parameter in global scope "
) ,
] ) )
2023-11-28 17:37:49 +08:00
} ;
match & func . node {
ExprKind ::Name { id , .. } if id = = & " TypeVar " . into ( ) = > {
2023-12-05 14:37:08 +08:00
let ExprKind ::Constant { value : Constant ::Str ( ty_name ) , .. } = & args [ 0 ] . node else {
2023-11-15 17:30:26 +08:00
return Err ( HashSet ::from ( [
format! ( " Expected string constant for first parameter of `TypeVar`, got {:?} " , & args [ 0 ] . node ) ,
] ) )
2023-12-05 14:37:08 +08:00
} ;
let generic_name : StrRef = ty_name . to_string ( ) . into ( ) ;
2022-04-18 16:02:48 +08:00
let constraints = args
. iter ( )
. skip ( 1 )
2023-11-15 17:30:26 +08:00
. map ( | x | -> Result < Type , HashSet < String > > {
2022-04-18 16:02:48 +08:00
let ty = parse_ast_to_type_annotation_kinds (
resolver ,
def_list ,
unifier ,
primitives ,
x ,
2023-12-11 14:41:58 +08:00
HashMap ::default ( ) ,
2023-12-05 14:37:08 +08:00
None ,
2022-04-18 16:02:48 +08:00
) ? ;
get_type_from_type_annotation_kinds (
2023-12-06 11:49:02 +08:00
def_list , unifier , & ty , & mut None
2022-04-18 16:02:48 +08:00
)
} )
. collect ::< Result < Vec < _ > , _ > > ( ) ? ;
2023-12-05 14:37:08 +08:00
let loc = func . location ;
if constraints . len ( ) = = 1 {
2023-11-15 17:30:26 +08:00
return Err ( HashSet ::from ( [
format! ( " A single constraint is not allowed (at {loc} ) " ) ,
] ) )
2023-12-05 14:37:08 +08:00
}
Ok ( unifier . get_fresh_var_with_range ( & constraints , Some ( generic_name ) , Some ( loc ) ) . 0 )
}
ExprKind ::Name { id , .. } if id = = & " ConstGeneric " . into ( ) = > {
if args . len ( ) ! = 2 {
2023-11-15 17:30:26 +08:00
return Err ( HashSet ::from ( [
format! ( " Expected 2 arguments for `ConstGeneric`, got {} " , args . len ( ) ) ,
] ) )
2023-12-05 14:37:08 +08:00
}
let ExprKind ::Constant { value : Constant ::Str ( ty_name ) , .. } = & args [ 0 ] . node else {
2023-11-15 17:30:26 +08:00
return Err ( HashSet ::from ( [
format! (
" Expected string constant for first parameter of `ConstGeneric`, got {:?} " ,
& args [ 0 ] . node
) ,
] ) )
2023-12-05 14:37:08 +08:00
} ;
let generic_name : StrRef = ty_name . to_string ( ) . into ( ) ;
let ty = parse_ast_to_type_annotation_kinds (
resolver ,
def_list ,
unifier ,
primitives ,
& args [ 1 ] ,
2023-12-11 14:41:58 +08:00
HashMap ::default ( ) ,
2023-12-05 14:37:08 +08:00
None ,
) ? ;
let constraint = get_type_from_type_annotation_kinds (
2023-12-06 11:49:02 +08:00
def_list , unifier , & ty , & mut None
2023-12-05 14:37:08 +08:00
) ? ;
let loc = func . location ;
Ok ( unifier . get_fresh_const_generic_var ( constraint , Some ( generic_name ) , Some ( loc ) ) . 0 )
2022-04-18 16:02:48 +08:00
}
2023-11-28 17:37:49 +08:00
2023-11-15 17:30:26 +08:00
_ = > Err ( HashSet ::from ( [
format! (
" expression {var:?} cannot be handled as a generic parameter in global scope "
) ,
] ) )
2022-04-18 16:02:48 +08:00
}
}
fn handle_assignment_pattern (
targets : & [ Expr ] ,
value : & Expr ,
resolver : & ( dyn SymbolResolver + Send + Sync ) ,
internal_resolver : & ResolverInternal ,
def_list : & [ Arc < RwLock < TopLevelDef > > ] ,
unifier : & mut Unifier ,
primitives : & PrimitiveStore ,
) -> Result < ( ) , String > {
if targets . len ( ) = = 1 {
match & targets [ 0 ] . node {
ExprKind ::Name { id , .. } = > {
if let Ok ( var ) = handle_typevar_definition (
2023-11-26 09:09:24 +08:00
value ,
2022-04-18 16:02:48 +08:00
resolver ,
def_list ,
unifier ,
primitives ,
) {
internal_resolver . add_id_type ( * id , var ) ;
Ok ( ( ) )
} else if let Ok ( val ) =
2023-11-26 09:09:24 +08:00
parse_parameter_default_value ( value , resolver )
2022-04-18 16:02:48 +08:00
{
internal_resolver . add_module_global ( * id , val ) ;
Ok ( ( ) )
} else {
2023-11-28 17:37:49 +08:00
Err ( format! ( " fails to evaluate this expression ` {:?} ` as a constant or generic parameter at {} " ,
2022-04-18 16:02:48 +08:00
targets [ 0 ] . node ,
targets [ 0 ] . location ,
) )
}
}
ExprKind ::List { elts , .. } | ExprKind ::Tuple { elts , .. } = > {
handle_assignment_pattern (
elts ,
value ,
resolver ,
internal_resolver ,
def_list ,
unifier ,
primitives ,
) ? ;
Ok ( ( ) )
}
_ = > Err ( format! (
" assignment to {:?} is not supported at {} " ,
targets [ 0 ] , targets [ 0 ] . location
) ) ,
}
} else {
match & value . node {
ExprKind ::List { elts , .. } | ExprKind ::Tuple { elts , .. } = > {
2023-12-11 14:41:58 +08:00
if elts . len ( ) = = targets . len ( ) {
2022-04-18 16:02:48 +08:00
for ( tar , val ) in targets . iter ( ) . zip ( elts ) {
handle_assignment_pattern (
std ::slice ::from_ref ( tar ) ,
val ,
resolver ,
internal_resolver ,
def_list ,
unifier ,
primitives ,
) ? ;
}
Ok ( ( ) )
2023-12-11 14:41:58 +08:00
} else {
Err ( format! (
" number of elements to unpack does not match (expect {}, found {}) at {} " ,
targets . len ( ) ,
elts . len ( ) ,
value . location
) )
2022-04-18 16:02:48 +08:00
}
}
_ = > Err ( format! (
" unpack of this expression is not supported at {} " ,
value . location
) ) ,
}
}
}
2020-12-17 22:20:30 +08:00
fn main ( ) {
2023-09-06 16:34:08 +08:00
let cli = CommandLineArgs ::parse ( ) ;
2023-09-12 14:41:34 +08:00
let CommandLineArgs {
file_name ,
threads ,
opt_level ,
emit_llvm ,
triple ,
mcpu ,
target_features ,
} = cli ;
Target ::initialize_all ( & InitializationConfig ::default ( ) ) ;
let host_target_machine = CodeGenTargetMachineOptions ::from_host ( ) ;
let triple = triple . unwrap_or ( host_target_machine . triple . clone ( ) ) ;
let mcpu = mcpu
. map ( | arg | if arg = = " native " { host_target_machine . cpu . clone ( ) } else { arg } )
. unwrap_or_default ( ) ;
let target_features = target_features . unwrap_or_default ( ) ;
2023-10-13 11:28:30 +08:00
let threads = if is_multithreaded ( ) {
if threads = = 0 {
std ::thread ::available_parallelism ( )
. map ( | threads | threads . get ( ) as u32 )
. unwrap_or ( 1 u32 )
} else {
threads
}
2023-10-13 14:57:09 +08:00
} else {
2023-10-13 11:28:30 +08:00
if threads ! = 1 {
2023-12-11 14:41:58 +08:00
println! ( " Warning: Number of threads specified in command-line but multithreading is disabled in LLVM at build time! Defaulting to single-threaded compilation " ) ;
2023-10-13 11:28:30 +08:00
}
1
2023-10-13 14:57:09 +08:00
} ;
2023-09-06 18:49:58 +08:00
let opt_level = match opt_level {
0 = > OptimizationLevel ::None ,
1 = > OptimizationLevel ::Less ,
2 = > OptimizationLevel ::Default ,
// The default behavior for -O<n> where n>3 defaults to O3 for both Clang and GCC
_ = > OptimizationLevel ::Aggressive ,
} ;
2021-09-22 14:30:52 +08:00
2021-12-28 10:59:17 +08:00
let program = match fs ::read_to_string ( file_name . clone ( ) ) {
2020-12-17 22:20:30 +08:00
Ok ( program ) = > program ,
2021-08-19 15:30:52 +08:00
Err ( err ) = > {
2023-12-11 14:41:58 +08:00
println! ( " Cannot open input file: {err} " ) ;
2021-08-19 15:30:52 +08:00
return ;
}
2021-10-16 22:17:36 +08:00
} ;
2021-08-19 15:30:52 +08:00
2021-09-19 16:19:16 +08:00
let primitive : PrimitiveStore = TopLevelComposer ::make_primitives ( ) . 0 ;
2022-02-21 18:27:46 +08:00
let ( mut composer , builtins_def , builtins_ty ) =
2023-12-11 14:41:58 +08:00
TopLevelComposer ::new ( vec! [ ] , ComposerConfig ::default ( ) ) ;
2021-08-27 16:25:59 +08:00
2021-09-19 16:19:16 +08:00
let internal_resolver : Arc < ResolverInternal > = ResolverInternal {
id_to_type : builtins_ty . into ( ) ,
id_to_def : builtins_def . into ( ) ,
2023-12-11 14:41:58 +08:00
class_names : Mutex ::default ( ) ,
module_globals : Mutex ::default ( ) ,
str_store : Mutex ::default ( ) ,
2022-04-18 16:02:48 +08:00
} . into ( ) ;
2021-10-16 22:17:36 +08:00
let resolver =
Arc ::new ( Resolver ( internal_resolver . clone ( ) ) ) as Arc < dyn SymbolResolver + Send + Sync > ;
2021-09-22 14:45:56 +08:00
2021-12-28 10:59:17 +08:00
let parser_result = parser ::parse_program ( & program , file_name . into ( ) ) . unwrap ( ) ;
2021-08-27 16:25:59 +08:00
2023-12-11 14:41:58 +08:00
for stmt in parser_result {
2022-04-18 16:02:48 +08:00
match & stmt . node {
StmtKind ::Assign { targets , value , .. } = > {
let def_list = composer . extract_def_list ( ) ;
let unifier = & mut composer . unifier ;
let primitives = & composer . primitives_ty ;
if let Err ( err ) = handle_assignment_pattern (
targets ,
value ,
resolver . as_ref ( ) ,
internal_resolver . as_ref ( ) ,
& def_list ,
unifier ,
primitives ,
) {
2023-12-11 14:41:58 +08:00
eprintln! ( " {err} " ) ;
2022-04-18 16:02:48 +08:00
return ;
2021-12-01 22:44:53 +08:00
}
2022-04-18 16:02:48 +08:00
} ,
// allow (and ignore) "from __future__ import annotations"
StmtKind ::ImportFrom { module , names , .. }
if module = = & Some ( " __future__ " . into ( ) ) & & names . len ( ) = = 1 & & names [ 0 ] . name = = " annotations " . into ( ) = > ( ) ,
_ = > {
let ( name , def_id , ty ) =
2023-12-11 14:41:58 +08:00
composer . register_top_level ( stmt , Some ( resolver . clone ( ) ) , " __main__ " , true ) . unwrap ( ) ;
2022-04-18 16:02:48 +08:00
internal_resolver . add_id_def ( name , def_id ) ;
if let Some ( ty ) = ty {
internal_resolver . add_id_type ( name , ty ) ;
2021-11-23 07:32:09 +08:00
}
}
2021-08-27 16:25:59 +08:00
}
2021-08-27 11:39:36 +08:00
}
2021-08-27 11:13:43 +08:00
2022-02-21 18:27:46 +08:00
let signature = FunSignature { args : vec ! [ ] , ret : primitive . int32 , vars : HashMap ::new ( ) } ;
2021-10-17 13:02:18 +08:00
let mut store = ConcreteTypeStore ::new ( ) ;
let mut cache = HashMap ::new ( ) ;
let signature = store . from_signature ( & mut composer . unifier , & primitive , & signature , & mut cache ) ;
let signature = store . add_cty ( signature ) ;
2021-09-19 16:19:16 +08:00
composer . start_analysis ( true ) . unwrap ( ) ;
2021-08-27 16:25:59 +08:00
2021-09-19 16:19:16 +08:00
let top_level = Arc ::new ( composer . make_top_level_context ( ) ) ;
2021-08-25 15:30:36 +08:00
2021-08-27 16:25:59 +08:00
let instance = {
let defs = top_level . definitions . read ( ) ;
2022-02-21 18:27:46 +08:00
let mut instance = defs [ resolver
. get_identifier_def ( " run " . into ( ) )
. unwrap_or_else ( | _ | panic! ( " cannot find run() entry point " ) )
. 0 ]
. write ( ) ;
if let TopLevelDef ::Function { instance_to_stmt , instance_to_symbol , .. } = & mut * instance {
2023-12-11 14:41:58 +08:00
instance_to_symbol . insert ( String ::new ( ) , " run " . to_string ( ) ) ;
2021-08-27 16:25:59 +08:00
instance_to_stmt [ " " ] . clone ( )
} else {
unreachable! ( )
}
} ;
2021-09-19 16:19:16 +08:00
2023-09-06 17:50:17 +08:00
let llvm_options = CodeGenLLVMOptions {
2023-09-06 18:49:58 +08:00
opt_level ,
2023-09-12 14:41:34 +08:00
target : CodeGenTargetMachineOptions {
triple ,
cpu : mcpu ,
features : target_features ,
2023-12-04 19:09:50 +08:00
reloc_mode : RelocMode ::PIC ,
2023-09-12 14:41:34 +08:00
.. host_target_machine
} ,
2023-09-06 17:50:17 +08:00
} ;
2023-09-11 13:12:46 +08:00
2021-08-19 15:30:52 +08:00
let task = CodeGenTask {
2023-12-11 14:41:58 +08:00
subst : Vec ::default ( ) ,
2021-08-19 15:30:52 +08:00
symbol_name : " run " . to_string ( ) ,
2021-08-27 16:25:59 +08:00
body : instance . body ,
2021-08-19 15:30:52 +08:00
signature ,
2021-08-27 16:25:59 +08:00
resolver ,
2021-10-17 13:02:18 +08:00
store ,
unifier_index : instance . unifier_id ,
2021-08-27 16:25:59 +08:00
calls : instance . calls ,
2021-11-20 19:50:25 +08:00
id : 0 ,
2021-08-19 15:30:52 +08:00
} ;
2021-08-27 16:25:59 +08:00
2023-12-11 14:41:58 +08:00
let membuffers : Arc < Mutex < Vec < Vec < u8 > > > > = Arc ::default ( ) ;
2022-01-09 02:11:33 +08:00
let membuffer = membuffers . clone ( ) ;
2022-02-13 17:21:42 +08:00
2022-01-09 02:11:33 +08:00
let f = Arc ::new ( WithCall ::new ( Box ::new ( move | module | {
let buffer = module . write_bitcode_to_memory ( ) ;
let buffer = buffer . as_slice ( ) . into ( ) ;
membuffer . lock ( ) . push ( buffer ) ;
2021-08-19 15:30:52 +08:00
} ) ) ) ;
2021-10-16 22:17:36 +08:00
let threads = ( 0 .. threads )
2023-12-11 14:41:58 +08:00
. map ( | i | Box ::new ( DefaultCodeGenerator ::new ( format! ( " module {i} " ) , 64 ) ) )
2021-10-16 22:17:36 +08:00
. collect ( ) ;
2023-12-08 17:43:32 +08:00
let ( registry , handles ) = WorkerRegistry ::create_workers ( threads , top_level , & llvm_options , & f ) ;
2021-08-19 15:30:52 +08:00
registry . add_task ( task ) ;
registry . wait_tasks_complete ( handles ) ;
2021-09-22 14:30:52 +08:00
2022-01-09 02:11:33 +08:00
let buffers = membuffers . lock ( ) ;
let context = inkwell ::context ::Context ::create ( ) ;
let main = context
. create_module_from_ir ( MemoryBuffer ::create_from_memory_range ( & buffers [ 0 ] , " main " ) )
. unwrap ( ) ;
2023-09-15 13:26:15 +08:00
if emit_llvm {
main . write_bitcode_to_path ( Path ::new ( " main.bc " ) ) ;
}
for ( idx , buffer ) in buffers . iter ( ) . skip ( 1 ) . enumerate ( ) {
2022-01-09 02:11:33 +08:00
let other = context
. create_module_from_ir ( MemoryBuffer ::create_from_memory_range ( buffer , " main " ) )
. unwrap ( ) ;
2023-09-15 13:26:15 +08:00
if emit_llvm {
2023-12-11 14:41:58 +08:00
other . write_bitcode_to_path ( Path ::new ( & format! ( " module {idx} .bc " ) ) ) ;
2023-09-15 13:26:15 +08:00
}
2022-01-09 02:11:33 +08:00
main . link_in_module ( other ) . unwrap ( ) ;
}
2023-09-29 16:41:32 +08:00
let irrt = load_irrt ( & context ) ;
if emit_llvm {
irrt . write_bitcode_to_path ( Path ::new ( " irrt.bc " ) ) ;
}
main . link_in_module ( irrt ) . unwrap ( ) ;
2022-01-09 02:11:33 +08:00
let mut function_iter = main . get_first_function ( ) ;
while let Some ( func ) = function_iter {
if func . count_basic_blocks ( ) > 0 & & func . get_name ( ) . to_str ( ) . unwrap ( ) ! = " run " {
func . set_linkage ( inkwell ::module ::Linkage ::Private ) ;
}
function_iter = func . get_next_function ( ) ;
}
2023-09-11 13:12:46 +08:00
let target_machine = llvm_options . target
. create_target_machine ( llvm_options . opt_level )
2022-01-09 02:11:33 +08:00
. expect ( " couldn't create target machine " ) ;
2023-09-11 13:14:35 +08:00
let pass_options = PassBuilderOptions ::create ( ) ;
pass_options . set_merge_functions ( true ) ;
2023-09-21 15:34:05 +08:00
let passes = format! ( " default<O {} > " , opt_level as u32 ) ;
let result = main . run_passes ( passes . as_str ( ) , & target_machine , pass_options ) ;
2023-09-11 13:14:35 +08:00
if let Err ( err ) = result {
panic! ( " Failed to run optimization for module `main`: {} " , err . to_string ( ) ) ;
}
2022-01-09 02:11:33 +08:00
target_machine
2022-02-21 18:27:46 +08:00
. write_to_file ( & main , FileType ::Object , Path ::new ( " module.o " ) )
2022-01-09 02:11:33 +08:00
. expect ( " couldn't write module to file " ) ;
2020-12-17 22:20:30 +08:00
}