mod circuit;
pub use circuit::*;
mod function;
pub use function::*;
use crate::{AsgConvertError, ImportResolver, InnerScope, Input, Scope};
use leo_ast::{Identifier, Package, PackageAccess, Span};
use indexmap::IndexMap;
use std::{cell::RefCell, sync::Arc};
use uuid::Uuid;
#[derive(Debug, Clone)]
pub struct InternalProgram {
pub id: Uuid,
pub name: String,
pub imported_modules: IndexMap<String, Program>,
pub test_functions: IndexMap<String, (Arc<FunctionBody>, Option<Identifier>)>,
pub functions: IndexMap<String, Arc<FunctionBody>>,
pub circuits: IndexMap<String, Arc<CircuitBody>>,
pub scope: Scope,
}
pub type Program = Arc<RefCell<InternalProgram>>;
enum ImportSymbol {
Direct(String),
Alias(String, String),
All,
}
fn resolve_import_package(
output: &mut Vec<(Vec<String>, ImportSymbol, Span)>,
mut package_segments: Vec<String>,
package: &Package,
) {
package_segments.push(package.name.name.clone());
resolve_import_package_access(output, package_segments, &package.access);
}
fn resolve_import_package_access(
output: &mut Vec<(Vec<String>, ImportSymbol, Span)>,
package_segments: Vec<String>,
package: &PackageAccess,
) {
match package {
PackageAccess::Star(span) => {
output.push((package_segments, ImportSymbol::All, span.clone()));
}
PackageAccess::SubPackage(subpackage) => {
resolve_import_package(output, package_segments, &*subpackage);
}
PackageAccess::Symbol(symbol) => {
let span = symbol.symbol.span.clone();
let symbol = if let Some(alias) = symbol.alias.as_ref() {
ImportSymbol::Alias(symbol.symbol.name.clone(), alias.name.clone())
} else {
ImportSymbol::Direct(symbol.symbol.name.clone())
};
output.push((package_segments, symbol, span));
}
PackageAccess::Multiple(subaccesses) => {
for subaccess in subaccesses.iter() {
resolve_import_package_access(output, package_segments.clone(), &subaccess);
}
}
}
}
impl InternalProgram {
pub fn new<T: ImportResolver + 'static>(
program: &leo_ast::Program,
import_resolver: &mut T,
) -> Result<Program, AsgConvertError> {
let mut imported_symbols: Vec<(Vec<String>, ImportSymbol, Span)> = vec![];
for import in program.imports.iter() {
resolve_import_package(&mut imported_symbols, vec![], &import.package);
}
let mut deduplicated_imports: IndexMap<Vec<String>, Span> = IndexMap::new();
for (package, _symbol, span) in imported_symbols.iter() {
deduplicated_imports.insert(package.clone(), span.clone());
}
let mut wrapped_resolver = crate::CoreImportResolver(import_resolver);
let mut resolved_packages: IndexMap<Vec<String>, Program> = IndexMap::new();
for (package, span) in deduplicated_imports.iter() {
let pretty_package = package.join(".");
let resolved_package =
match wrapped_resolver.resolve_package(&package.iter().map(|x| &**x).collect::<Vec<_>>()[..], span)? {
Some(x) => x,
None => return Err(AsgConvertError::unresolved_import(&*pretty_package, &Span::default())),
};
resolved_packages.insert(package.clone(), resolved_package);
}
let mut imported_functions: IndexMap<String, Arc<FunctionBody>> = IndexMap::new();
let mut imported_circuits: IndexMap<String, Arc<CircuitBody>> = IndexMap::new();
for (package, symbol, span) in imported_symbols.into_iter() {
let pretty_package = package.join(".");
let resolved_package = resolved_packages
.get(&package)
.expect("could not find preloaded package");
let resolved_package = resolved_package.borrow();
match symbol {
ImportSymbol::All => {
imported_functions.extend(resolved_package.functions.clone().into_iter());
imported_circuits.extend(resolved_package.circuits.clone().into_iter());
}
ImportSymbol::Direct(name) => {
if let Some(function) = resolved_package.functions.get(&name) {
imported_functions.insert(name.clone(), function.clone());
} else if let Some(function) = resolved_package.circuits.get(&name) {
imported_circuits.insert(name.clone(), function.clone());
} else {
return Err(AsgConvertError::unresolved_import(
&*format!("{}.{}", pretty_package, name),
&span,
));
}
}
ImportSymbol::Alias(name, alias) => {
if let Some(function) = resolved_package.functions.get(&name) {
imported_functions.insert(alias.clone(), function.clone());
} else if let Some(function) = resolved_package.circuits.get(&name) {
imported_circuits.insert(alias.clone(), function.clone());
} else {
return Err(AsgConvertError::unresolved_import(
&*format!("{}.{}", pretty_package, name),
&span,
));
}
}
}
}
let import_scope = Arc::new(RefCell::new(InnerScope {
id: uuid::Uuid::new_v4(),
parent_scope: None,
circuit_self: None,
variables: IndexMap::new(),
functions: imported_functions
.iter()
.map(|(name, func)| (name.clone(), func.function.clone()))
.collect(),
circuits: imported_circuits
.iter()
.map(|(name, circuit)| (name.clone(), circuit.circuit.clone()))
.collect(),
function: None,
input: None,
}));
let mut proto_circuits = IndexMap::new();
for (name, circuit) in program.circuits.iter() {
assert_eq!(name.name, circuit.circuit_name.name);
let asg_circuit = Circuit::init(circuit);
proto_circuits.insert(name.name.clone(), asg_circuit);
}
let scope = Arc::new(RefCell::new(InnerScope {
input: Some(Input::new(&import_scope)),
id: uuid::Uuid::new_v4(),
parent_scope: Some(import_scope),
circuit_self: None,
variables: IndexMap::new(),
functions: IndexMap::new(),
circuits: proto_circuits
.iter()
.map(|(name, circuit)| (name.clone(), circuit.clone()))
.collect(),
function: None,
}));
for (name, circuit) in program.circuits.iter() {
assert_eq!(name.name, circuit.circuit_name.name);
let asg_circuit = proto_circuits.get(&name.name).unwrap();
asg_circuit.clone().from_ast(&scope, &circuit)?;
}
let mut proto_test_functions = IndexMap::new();
for (name, test_function) in program.tests.iter() {
assert_eq!(name.name, test_function.function.identifier.name);
let function = Arc::new(Function::from_ast(&scope, &test_function.function)?);
proto_test_functions.insert(name.name.clone(), function);
}
let mut proto_functions = IndexMap::new();
for (name, function) in program.functions.iter() {
assert_eq!(name.name, function.identifier.name);
let asg_function = Arc::new(Function::from_ast(&scope, function)?);
scope
.borrow_mut()
.functions
.insert(name.name.clone(), asg_function.clone());
proto_functions.insert(name.name.clone(), asg_function);
}
let mut test_functions = IndexMap::new();
for (name, test_function) in program.tests.iter() {
assert_eq!(name.name, test_function.function.identifier.name);
let function = proto_test_functions.get(&name.name).unwrap();
let body = Arc::new(FunctionBody::from_ast(
&scope,
&test_function.function,
function.clone(),
)?);
function.body.replace(Arc::downgrade(&body));
test_functions.insert(name.name.clone(), (body, test_function.input_file.clone()));
}
let mut functions = IndexMap::new();
for (name, function) in program.functions.iter() {
assert_eq!(name.name, function.identifier.name);
let asg_function = proto_functions.get(&name.name).unwrap();
let body = Arc::new(FunctionBody::from_ast(&scope, function, asg_function.clone())?);
asg_function.body.replace(Arc::downgrade(&body));
functions.insert(name.name.clone(), body);
}
let mut circuits = IndexMap::new();
for (name, circuit) in program.circuits.iter() {
assert_eq!(name.name, circuit.circuit_name.name);
let asg_circuit = proto_circuits.get(&name.name).unwrap();
let body = Arc::new(CircuitBody::from_ast(&scope, circuit, asg_circuit.clone())?);
asg_circuit.body.replace(Arc::downgrade(&body));
circuits.insert(name.name.clone(), body);
}
Ok(Arc::new(RefCell::new(InternalProgram {
id: Uuid::new_v4(),
name: program.name.clone(),
test_functions,
functions,
circuits,
imported_modules: resolved_packages
.into_iter()
.map(|(package, program)| (package.join("."), program))
.collect(),
scope,
})))
}
pub(crate) fn set_core_mapping(&self, mapping: &str) {
for (_, circuit) in self.circuits.iter() {
circuit.circuit.core_mapping.replace(Some(mapping.to_string()));
}
}
}
struct InternalIdentifierGenerator {
next: usize,
}
impl Iterator for InternalIdentifierGenerator {
type Item = String;
fn next(&mut self) -> Option<String> {
let out = format!("$_{}_", self.next);
self.next += 1;
Some(out)
}
}
pub fn reform_ast(program: &Program) -> leo_ast::Program {
let mut all_programs: IndexMap<String, Program> = IndexMap::new();
let mut program_stack = program.borrow().imported_modules.clone();
while let Some((module, program)) = program_stack.pop() {
if all_programs.contains_key(&module) {
continue;
}
all_programs.insert(module, program.clone());
program_stack.extend(program.borrow().imported_modules.clone());
}
all_programs.insert("".to_string(), program.clone());
let core_programs: Vec<_> = all_programs
.iter()
.filter(|(module, _)| module.starts_with("core."))
.map(|(module, program)| (module.clone(), program.clone()))
.collect();
all_programs.retain(|module, _| !module.starts_with("core."));
let mut all_circuits: IndexMap<String, Arc<CircuitBody>> = IndexMap::new();
let mut all_functions: IndexMap<String, Arc<FunctionBody>> = IndexMap::new();
let mut all_test_functions: IndexMap<String, (Arc<FunctionBody>, Option<Identifier>)> = IndexMap::new();
let mut identifiers = InternalIdentifierGenerator { next: 0 };
for (_, program) in all_programs.into_iter() {
let program = program.borrow();
for (name, circuit) in program.circuits.iter() {
let identifier = format!("{}{}", identifiers.next().unwrap(), name);
circuit.circuit.name.borrow_mut().name = identifier.clone();
all_circuits.insert(identifier, circuit.clone());
}
for (name, function) in program.functions.iter() {
let identifier = if name == "main" {
"main".to_string()
} else {
format!("{}{}", identifiers.next().unwrap(), name)
};
function.function.name.borrow_mut().name = identifier.clone();
all_functions.insert(identifier, function.clone());
}
for (name, function) in program.test_functions.iter() {
let identifier = format!("{}{}", identifiers.next().unwrap(), name);
function.0.function.name.borrow_mut().name = identifier.clone();
all_test_functions.insert(identifier, function.clone());
}
}
leo_ast::Program {
name: "ast_aggregate".to_string(),
imports: core_programs
.iter()
.map(|(module, _)| leo_ast::ImportStatement {
package: leo_ast::Package {
name: Identifier::new(module.clone()),
access: leo_ast::PackageAccess::Star(Span::default()),
span: Default::default(),
},
span: Span::default(),
})
.collect(),
expected_input: vec![],
tests: all_test_functions
.into_iter()
.map(|(_, (function, ident))| {
(function.function.name.borrow().clone(), leo_ast::TestFunction {
function: function.function.as_ref().into(),
input_file: ident,
})
})
.collect(),
functions: all_functions
.into_iter()
.map(|(_, function)| {
(
function.function.name.borrow().clone(),
function.function.as_ref().into(),
)
})
.collect(),
circuits: all_circuits
.into_iter()
.map(|(_, circuit)| (circuit.circuit.name.borrow().clone(), circuit.circuit.as_ref().into()))
.collect(),
}
}
impl Into<leo_ast::Program> for &InternalProgram {
fn into(self) -> leo_ast::Program {
leo_ast::Program {
name: self.name.clone(),
imports: vec![],
expected_input: vec![],
circuits: self
.circuits
.iter()
.map(|(_, circuit)| (circuit.circuit.name.borrow().clone(), circuit.circuit.as_ref().into()))
.collect(),
functions: self
.functions
.iter()
.map(|(_, function)| {
(
function.function.name.borrow().clone(),
function.function.as_ref().into(),
)
})
.collect(),
tests: self
.test_functions
.iter()
.map(|(_, function)| {
(function.0.function.name.borrow().clone(), leo_ast::TestFunction {
function: function.0.function.as_ref().into(),
input_file: function.1.clone(),
})
})
.collect(),
}
}
}