Skip to content

Commit

Permalink
Allow any expression as constants
Browse files Browse the repository at this point in the history
  This is only a start. It compiles, but with a few TODOs left open. In particular, it doesn't currently handle constants depending on other constants or functions; nor does it hoist constants.
  • Loading branch information
KtorZ committed Aug 2, 2024
1 parent ea032c9 commit 86d60af
Show file tree
Hide file tree
Showing 13 changed files with 174 additions and 254 deletions.
73 changes: 17 additions & 56 deletions crates/aiken-lang/src/ast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,7 @@ impl TypedModule {
pub fn register_definitions(
&self,
functions: &mut IndexMap<FunctionAccessKey, TypedFunction>,
constants: &mut IndexMap<FunctionAccessKey, TypedExpr>,
data_types: &mut IndexMap<DataTypeKey, TypedDataType>,
) {
for def in self.definitions() {
Expand Down Expand Up @@ -182,7 +183,17 @@ impl TypedModule {
}
}

Definition::TypeAlias(_) | Definition::ModuleConstant(_) | Definition::Use(_) => {}
Definition::ModuleConstant(ModuleConstant { name, value, .. }) => {
constants.insert(
FunctionAccessKey {
module_name: self.name.clone(),
function_name: name.clone(),
},
value.clone(),
);
}

Definition::TypeAlias(_) | Definition::Use(_) => {}
}
}
}
Expand Down Expand Up @@ -504,18 +515,17 @@ pub struct Use<PackageName> {
pub unqualified: Vec<UnqualifiedImport>,
}

pub type TypedModuleConstant = ModuleConstant<Rc<Type>>;
pub type UntypedModuleConstant = ModuleConstant<()>;
pub type TypedModuleConstant = ModuleConstant<TypedExpr>;
pub type UntypedModuleConstant = ModuleConstant<UntypedExpr>;

#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ModuleConstant<T> {
pub struct ModuleConstant<Expr> {
pub doc: Option<String>,
pub location: Span,
pub public: bool,
pub name: String,
pub annotation: Option<Annotation>,
pub value: Box<Constant>,
pub tipo: T,
pub value: Expr,
}

pub type TypedValidator = Validator<Rc<Type>, TypedArg, TypedExpr>;
Expand Down Expand Up @@ -589,7 +599,7 @@ pub enum Definition<T, Arg, Expr, PackageName> {

Use(Use<PackageName>),

ModuleConstant(ModuleConstant<T>),
ModuleConstant(ModuleConstant<Expr>),

Test(Function<T, Expr, ArgVia<Arg, Expr>>),

Expand Down Expand Up @@ -686,55 +696,6 @@ pub struct DefinitionLocation<'module> {
pub span: Span,
}

#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum Constant {
Int {
location: Span,
value: String,
base: Base,
},

String {
location: Span,
value: String,
},

ByteArray {
location: Span,
bytes: Vec<u8>,
preferred_format: ByteArrayFormatPreference,
},

CurvePoint {
location: Span,
point: Box<Curve>,
preferred_format: ByteArrayFormatPreference,
},
}

impl Constant {
pub fn tipo(&self) -> Rc<Type> {
match self {
Constant::Int { .. } => builtins::int(),
Constant::String { .. } => builtins::string(),
Constant::ByteArray { .. } => builtins::byte_array(),
Constant::CurvePoint { point, .. } => match point.as_ref() {
Curve::Bls12_381(Bls12_381Point::G1(_)) => builtins::g1_element(),
Curve::Bls12_381(Bls12_381Point::G2(_)) => builtins::g2_element(),
},
}
}

pub fn location(&self) -> Span {
match self {
Constant::Int { location, .. }
| Constant::String { location, .. }
| Constant::ByteArray { location, .. }
| Constant::CurvePoint { location, .. } => *location,
}
}
}

pub type TypedCallArg = CallArg<TypedExpr>;
pub type ParsedCallArg = CallArg<Option<UntypedExpr>>;

Expand Down
34 changes: 10 additions & 24 deletions crates/aiken-lang/src/format.rs
Original file line number Diff line number Diff line change
@@ -1,15 +1,15 @@
use crate::{
ast::{
Annotation, ArgBy, ArgName, ArgVia, AssignmentKind, AssignmentPattern, BinOp,
ByteArrayFormatPreference, CallArg, Constant, CurveType, DataType, Definition, Function,
ByteArrayFormatPreference, CallArg, CurveType, DataType, Definition, Function,
LogicalOpChainKind, ModuleConstant, OnTestFailure, Pattern, RecordConstructor,
RecordConstructorArg, RecordUpdateSpread, Span, TraceKind, TypeAlias, TypedArg, UnOp,
UnqualifiedImport, UntypedArg, UntypedArgVia, UntypedAssignmentKind, UntypedClause,
UntypedDefinition, UntypedFunction, UntypedIfBranch, UntypedModule, UntypedPattern,
UntypedRecordUpdateArg, Use, Validator, CAPTURE_VARIABLE,
},
docvec,
expr::{FnStyle, UntypedExpr, DEFAULT_ERROR_STR, DEFAULT_TODO_STR},
expr::{FnStyle, TypedExpr, UntypedExpr, DEFAULT_ERROR_STR, DEFAULT_TODO_STR},
parser::{
extra::{Comment, ModuleExtra},
token::Base,
Expand Down Expand Up @@ -329,34 +329,20 @@ impl<'comments> Formatter<'comments> {
})
}

fn const_expr<'a>(&mut self, value: &'a Constant) -> Document<'a> {
match value {
Constant::ByteArray {
bytes,
preferred_format,
..
} => self.bytearray(bytes, None, preferred_format),
Constant::CurvePoint {
point,
preferred_format,
..
} => self.bytearray(
&point.compress(),
Some(point.as_ref().into()),
preferred_format,
),
Constant::Int { value, base, .. } => self.int(value, base),
Constant::String { value, .. } => self.string(value),
}
fn const_expr<'a>(&mut self, _value: &'a UntypedExpr) -> Document<'a> {
todo!(
"format const_expr: surround complex expressions with a block, and leave simple expression without"
);
}

pub fn docs_const_expr<'a>(&mut self, name: &'a str, value: &'a Constant) -> Document<'a> {
pub fn docs_const_expr<'a>(&mut self, name: &'a str, value: &'a TypedExpr) -> Document<'a> {
let mut printer = tipo::pretty::Printer::new();
name.to_doc()
.append(": ")
.append(printer.print(&value.tipo()))
.append(" = ")
.append(self.const_expr(value))
// TODO: Show full expression in docs when simple enough
// .append(" = ")
// .append(self.const_expr(value))
}

fn documented_definition<'a>(&mut self, s: &'a UntypedDefinition) -> Document<'a> {
Expand Down
76 changes: 63 additions & 13 deletions crates/aiken-lang/src/gen_uplc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,9 @@ pub mod tree;
use self::{
air::Air,
builder::{
cast_validator_args, constants_ir, convert_type_to_data, extract_constant,
modify_cyclic_calls, modify_self_calls, rearrange_list_clauses, AssignmentProperties,
ClauseProperties, CodeGenSpecialFuncs, CycleFunctionNames, HoistableFunction, Variant,
cast_validator_args, convert_type_to_data, extract_constant, modify_cyclic_calls,
modify_self_calls, rearrange_list_clauses, AssignmentProperties, ClauseProperties,
CodeGenSpecialFuncs, CycleFunctionNames, HoistableFunction, Variant,
},
tree::{AirTree, TreePath},
};
Expand Down Expand Up @@ -57,6 +57,7 @@ pub struct CodeGenerator<'a> {
plutus_version: PlutusVersion,
/// immutable index maps
functions: IndexMap<&'a FunctionAccessKey, &'a TypedFunction>,
constants: IndexMap<&'a FunctionAccessKey, &'a TypedExpr>,
data_types: IndexMap<&'a DataTypeKey, &'a TypedDataType>,
module_types: IndexMap<&'a str, &'a TypeInfo>,
module_src: IndexMap<&'a str, &'a (String, LineNumbers)>,
Expand All @@ -81,6 +82,7 @@ impl<'a> CodeGenerator<'a> {
pub fn new(
plutus_version: PlutusVersion,
functions: IndexMap<&'a FunctionAccessKey, &'a TypedFunction>,
constants: IndexMap<&'a FunctionAccessKey, &'a TypedExpr>,
data_types: IndexMap<&'a DataTypeKey, &'a TypedDataType>,
module_types: IndexMap<&'a str, &'a TypeInfo>,
module_src: IndexMap<&'a str, &'a (String, LineNumbers)>,
Expand All @@ -89,6 +91,7 @@ impl<'a> CodeGenerator<'a> {
CodeGenerator {
plutus_version,
functions,
constants,
data_types,
module_types,
module_src,
Expand Down Expand Up @@ -319,12 +322,7 @@ impl<'a> CodeGenerator<'a> {

TypedExpr::Var {
constructor, name, ..
} => match &constructor.variant {
ValueConstructorVariant::ModuleConstant { literal, .. } => {
constants_ir(literal)
}
_ => AirTree::var(constructor.clone(), name, ""),
},
} => AirTree::var(constructor.clone(), name, ""),

TypedExpr::Fn { args, body, .. } => AirTree::anon_func(
args.iter()
Expand Down Expand Up @@ -784,8 +782,16 @@ impl<'a> CodeGenerator<'a> {
AirTree::builtin(*builtin, tipo.clone(), vec![])
}
}
ModuleValueConstructor::Constant { literal, .. } => {
builder::constants_ir(literal)
ModuleValueConstructor::Constant { module, name, .. } => {
let type_info = self.module_types.get(module_name.as_str()).unwrap();

let value = type_info.values.get(name).unwrap();

AirTree::var(
ValueConstructor::public(tipo.clone(), value.variant.clone()),
format!("{module}_{name}"),
"",
)
}
},

Expand Down Expand Up @@ -4057,8 +4063,52 @@ impl<'a> CodeGenerator<'a> {
}
.into(),
)),
ValueConstructorVariant::ModuleConstant { .. } => {
unreachable!("{:#?}, {}", constructor, name)
ValueConstructorVariant::ModuleConstant { module, name, .. } => {
let access_key = FunctionAccessKey {
module_name: module.clone(),
function_name: name.clone(),
};

let definition = self
.constants
.get(&access_key)
.unwrap_or_else(|| panic!("unknown constant {module}.{name}"));

let mut value =
AirTree::no_op(self.build(&definition, &access_key.module_name, &[]));

value.traverse_tree_with(
&mut |air_tree, _| {
erase_opaque_type_operations(air_tree, &self.data_types);
},
true,
);

let term = self
.uplc_code_gen(value.to_vec())
.constr_fields_exposer()
.constr_index_exposer();

let mut program: Program<Name> = Program {
version: (1, 0, 0),
term: self.special_functions.apply_used_functions(term),
};

let mut interner = CodeGenInterner::new();

interner.program(&mut program);

let eval_program: Program<NamedDeBruijn> =
program.remove_no_inlines().try_into().unwrap();

Some(
eval_program
.eval(ExBudget::max())
.result()
.unwrap_or_else(|e| panic!("Failed to evaluate constant: {e:#?}"))
.try_into()
.unwrap(),
)
}
ValueConstructorVariant::ModuleFn {
name: func_name,
Expand Down
13 changes: 2 additions & 11 deletions crates/aiken-lang/src/gen_uplc/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@ use super::{
};
use crate::{
ast::{
Constant, DataTypeKey, FunctionAccessKey, Pattern, Span, TraceLevel, TypedArg,
TypedAssignmentKind, TypedClause, TypedDataType, TypedPattern,
DataTypeKey, FunctionAccessKey, Pattern, Span, TraceLevel, TypedArg, TypedAssignmentKind,
TypedClause, TypedDataType, TypedPattern,
},
builtins::{data, function, int, list, void},
expr::TypedExpr,
Expand Down Expand Up @@ -288,15 +288,6 @@ impl Default for CodeGenSpecialFuncs {
}
}

pub fn constants_ir(literal: &Constant) -> AirTree {
match literal {
Constant::Int { value, .. } => AirTree::int(value),
Constant::String { value, .. } => AirTree::string(value),
Constant::ByteArray { bytes, .. } => AirTree::byte_array(bytes.clone()),
Constant::CurvePoint { point, .. } => AirTree::curve(*point.as_ref()),
}
}

pub fn get_generic_variant_name(t: &Rc<Type>) -> String {
let uplc_type = t.get_uplc_type();

Expand Down
Loading

0 comments on commit 86d60af

Please sign in to comment.