Skip to content

Commit

Permalink
feat(macros): add sub mul div zero derives (#328)
Browse files Browse the repository at this point in the history
<!--- Please provide a general summary of your changes in the title
above -->

## Pull Request type

<!-- Please try to limit your pull request to one type; submit multiple
pull requests if needed. -->

Please check the type of change your PR introduces:

- [ ] Bugfix
- [x] Feature
- [ ] Code style update (formatting, renaming)
- [ ] Refactoring (no functional changes, no API changes)
- [ ] Build-related changes
- [ ] Documentation content changes
- [ ] Other (please describe):


## What is the new behavior?

PR introduces a bunch of new macros that allow to derive `Add`, `Sub`,
`Mul`, `Div`, `AddAssign`, `SubAssign`, `MulAssign`, `DivAssign` and
`Zero` traits. Now we can have a struct like this:

```
#[derive(Add, AddAssign, Sub, SubAssign, Mul, MulAssign, Div, DivAssign, Zero)]
struct Foo {
    x: u8,
    y: u16,
    z: u32
}
```

and all the basic ops are available out of the box.

The macros work for more complex types as well (generic & nested
structs).

## Does this introduce a breaking change?

- [ ] Yes
- [x] No

## Other information

<!-- Any other information that is important to this PR, such as
screenshots of how the component looks before and after the change. -->

I've cleaned up the package structure to be more modular as it grows.
  • Loading branch information
milancermak authored Aug 30, 2024
1 parent 7c19379 commit fdf4098
Show file tree
Hide file tree
Showing 10 changed files with 730 additions and 68 deletions.
62 changes: 60 additions & 2 deletions packages/macros/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,70 @@

A collection of useful macros.

- [`pow!`](#pow)
- [Add, Sub, Mul, Div derives](#add-sub-mul-div-derives)
- [AddAssign, SubAssign, MulAssign, DivAssign derives](#addassign-subassign-mulassign-divassign-derives)
- [pow!](#pow)
- [Zero derive](#zero-derive)

## `pow!`
## Add, Sub, Mul, Div derives

These macros add the ability to use addition `+`, subtraction `-`, multiplication `*` and division `/` operators on the type that derives them. All individual members of the struct must already support a particular operation for a derive to work.

```rust
#[derive(Add, Sub, Mul, Div)]
struct Point {
x: u32,
y: u32
}

let a = Point { x: 6, y: 32 };
let b = Point { x: 2, y: 2 };
let c = a + b; // Point { x: 8, y: 34 }
let d = a - b; // Point { x: 4, y: 30 }
let e = a * b; // Point { x: 12, y: 64 }
let f = a / b; // Point { x: 3, y: 16 }
```

## AddAssign, SubAssign, MulAssign, DivAssign derives

These macros add the ability to use add and assign `+=`, subtract and assign `-=`, multiply and assign `*=` and divide and assign `/=` operators on the type that derives them. All individual members of the struct must already support the particular operation for a derive to work.

```rust
#[derive(AddAssign, SubAssign, MulAssign, DivAssign)]
struct Point {
x: u32,
y: u32
}

let mut a = Point { x: 6, y: 32 };
let b = Point { x: 2, y: 2 };
a += b; // Point { x: 8, y: 34 }
a -= b; // Point { x: 6, y: 32 }
a *= b; // Point { x: 12, y: 64 }
a /= b; // Point { x: 6, y: 32 };
```

## pow!

Power function. Takes two arguments, `x, y`, calculates the value of `x` raised to the power of `y`.

```cairo
const MEGABYTE: u64 = pow!(2, 20); // will be set to 1048576
```

## Zero derive

Adds implementation of the `core::num::traits::Zero` trait.

All members of the struct must already implement the `Zero` trait.

```rust
#[derive(Zero, PartialEq, Debug)]
struct Point {
x: u64,
y: u64,
}
assert_eq!(Point { x: 0, y: 0 }, Zero::zero());
assert!(Point { x: 0, y: 0 }.is_zero());
assert!(Point { x: 1, y: 0 }.is_non_zero());
```
57 changes: 4 additions & 53 deletions packages/macros/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,54 +1,5 @@
use bigdecimal::{num_traits::pow, BigDecimal};
use cairo_lang_macro::{inline_macro, Diagnostic, ProcMacroResult, TokenStream};
use cairo_lang_parser::utils::SimpleParserDatabase;
use cairo_lang_syntax::node::kind::SyntaxKind::Arg;
mod num_traits;
mod pow;
mod zero_trait;

/// Compile-time power function.
///
/// Takes two arguments, `x, y`, calculates the value of `x` raised to the power of `y`.
///
/// ```
/// const MEGABYTE: u64 = pow!(2, 20);
/// assert_eq!(MEGABYTE, 1048576);
/// ```
#[inline_macro]
pub fn pow(token_stream: TokenStream) -> ProcMacroResult {
let db = SimpleParserDatabase::default();
let (parsed, _diag) = db.parse_virtual_with_diagnostics(token_stream);

let macro_args: Vec<String> = parsed
.descendants(&db)
.filter_map(|node| {
if let Arg = node.kind(&db) {
Some(node.get_text(&db))
} else {
None
}
})
.collect();

if macro_args.len() != 2 {
return ProcMacroResult::new(TokenStream::empty())
.with_diagnostics(Diagnostic::error("Invalid number of arguments").into());
}

let base: BigDecimal = match macro_args[0].parse() {
Ok(val) => val,
Err(_) => {
return ProcMacroResult::new(TokenStream::empty())
.with_diagnostics(Diagnostic::error("Invalid base value").into());
}
};

let exp: usize = match macro_args[1].parse() {
Ok(val) => val,
Err(_) => {
return ProcMacroResult::new(TokenStream::empty())
.with_diagnostics(Diagnostic::error("Invalid exponent value").into());
}
};

let result: BigDecimal = pow(base, exp);

ProcMacroResult::new(TokenStream::new(result.to_string()))
}
mod parse;
222 changes: 222 additions & 0 deletions packages/macros/src/num_traits.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,222 @@
use cairo_lang_macro::{derive_macro, ProcMacroResult, TokenStream};

use crate::parse::{parse_struct_info, StructInfo};

struct OpInfo {
trait_name: String,
fn_name: String,
operator: String,
}

fn generate_op_trait_impl(op_info: &OpInfo, s: &StructInfo) -> String {
let generic_params = s
.generic_params
.as_ref()
.map_or(String::new(), |params| format!("<{}>", params.join(", ")));

let trait_bounds = s
.generic_params
.as_ref()
.map_or_else(String::new, |params| {
let bounds = params
.iter()
.flat_map(|param| {
vec![
format!("+core::traits::{}<{}>", op_info.trait_name, param),
format!("+core::traits::Drop<{}>", param),
]
})
.collect::<Vec<_>>()
.join(",\n");
format!("<{},\n{}>", params.join(", "), bounds)
});

let members_op = s
.members
.iter()
.map(|member| format!("{0}: lhs.{0} {1} rhs.{0}", member, op_info.operator))
.collect::<Vec<_>>()
.join(", ");

format!(
"\n
impl {0}{1}{2}
of core::traits::{1}<{0}{3}> {{
fn {4}(lhs: {0}{3}, rhs: {0}{3}) -> {0}{3} {{
{0} {{ {5} }}
}}
}}\n",
s.name, op_info.trait_name, trait_bounds, generic_params, op_info.fn_name, members_op
)
}

fn generate_op_assign_trait_impl(op_info: &OpInfo, s: &StructInfo) -> String {
let generic_params = s
.generic_params
.as_ref()
.map_or(String::new(), |params| format!("<{}>", params.join(", ")));

let trait_bounds = s
.generic_params
.as_ref()
.map_or_else(String::new, |params| {
let bounds = params
.iter()
.flat_map(|param| {
vec![
format!("+core::ops::{0}Assign<{1}, {1}>", op_info.trait_name, param),
format!("+core::traits::Drop<{}>", param),
]
})
.collect::<Vec<_>>()
.join(",\n");
format!("<{},\n{}>", params.join(", "), bounds)
});

let members_op = s
.members
.iter()
.map(|member| format!("self.{0} {1}= rhs.{0}", member, op_info.operator))
.collect::<Vec<_>>()
.join(";\n ");

format!(
"\n
impl {0}{1}Assign{2}
of core::ops::{1}Assign<{0}{3}, {0}{3}> {{
fn {4}_assign(ref self: {0}{3}, rhs: {0}{3}) {{
{5};
}}
}}\n",
s.name, op_info.trait_name, trait_bounds, generic_params, op_info.fn_name, members_op
)
}

/// Adds implementation for the `core::traits::Add` trait.
///
/// Allows you to use the `+` oprator on a type. All members of
/// the struct must already implement the `Add` trait.
#[derive_macro]
pub fn add(token_stream: TokenStream) -> ProcMacroResult {
let op = OpInfo {
trait_name: "Add".to_string(),
fn_name: "add".to_string(),
operator: "+".to_string(),
};

let s = parse_struct_info(token_stream);

ProcMacroResult::new(TokenStream::new(generate_op_trait_impl(&op, &s)))
}

/// Adds implementation for the `core::traits::Sub` trait.
///
/// Allows you to use the `-` operator on a type. All members of
/// the struct must already implement the `Sub` trait.
#[derive_macro]
pub fn sub(token_stream: TokenStream) -> ProcMacroResult {
let op = OpInfo {
trait_name: "Sub".to_string(),
fn_name: "sub".to_string(),
operator: "-".to_string(),
};
let s = parse_struct_info(token_stream);

ProcMacroResult::new(TokenStream::new(generate_op_trait_impl(&op, &s)))
}

/// Adds implementation for the `core::traits::Mul` trait.
///
/// Allows you to use the `*` operator on a type. All members of
/// the struct must already implement the `Mul` trait.
#[derive_macro]
pub fn mul(token_stream: TokenStream) -> ProcMacroResult {
let op = OpInfo {
trait_name: "Mul".to_string(),
fn_name: "mul".to_string(),
operator: "*".to_string(),
};
let s = parse_struct_info(token_stream);

ProcMacroResult::new(TokenStream::new(generate_op_trait_impl(&op, &s)))
}

/// Adds implementation for the `core::traits::Div` trait.
///
/// Allows you to use the `/` operator on a type. All members of
/// the struct must already implement the `Div` trait.
#[derive_macro]
pub fn div(token_stream: TokenStream) -> ProcMacroResult {
let op = OpInfo {
trait_name: "Div".to_string(),
fn_name: "div".to_string(),
operator: "/".to_string(),
};
let s = parse_struct_info(token_stream);

ProcMacroResult::new(TokenStream::new(generate_op_trait_impl(&op, &s)))
}

/// Adds implementation for the `core::ops::AddAssign` trait.
///
/// Allows you to use the `+=` operator on a type. All members of
/// the struct must already implement the `AddAssign` trait.
#[derive_macro]
fn add_assign(token_stream: TokenStream) -> ProcMacroResult {
let op = OpInfo {
trait_name: "Add".to_string(),
fn_name: "add".to_string(),
operator: "+".to_string(),
};
let s = parse_struct_info(token_stream);

ProcMacroResult::new(TokenStream::new(generate_op_assign_trait_impl(&op, &s)))
}

/// Adds implementation for the `core::ops::SubAssign` trait.
///
/// Allows you to use the `-=` operator on a type. All members of
/// the struct must already implement the `SubAssign` trait.
#[derive_macro]
fn sub_assign(token_stream: TokenStream) -> ProcMacroResult {
let op = OpInfo {
trait_name: "Sub".to_string(),
fn_name: "sub".to_string(),
operator: "-".to_string(),
};
let s = parse_struct_info(token_stream);

ProcMacroResult::new(TokenStream::new(generate_op_assign_trait_impl(&op, &s)))
}

/// Adds implementation for the `core::ops::MulAssign` trait.
///
/// Allows you to use the `*=` operator on a type. All members of
/// the struct must already implement the `MulAssign` trait.
#[derive_macro]
fn mul_assign(token_stream: TokenStream) -> ProcMacroResult {
let op = OpInfo {
trait_name: "Mul".to_string(),
fn_name: "mul".to_string(),
operator: "*".to_string(),
};
let s = parse_struct_info(token_stream);

ProcMacroResult::new(TokenStream::new(generate_op_assign_trait_impl(&op, &s)))
}

/// Adds implementation for the `core::ops::DivAssign` trait.
///
/// Allows you to use the `/=` operator on a type. All members of
/// the struct must already implement the `DivAssign` trait.
#[derive_macro]
fn div_assign(token_stream: TokenStream) -> ProcMacroResult {
let op = OpInfo {
trait_name: "Div".to_string(),
fn_name: "div".to_string(),
operator: "/".to_string(),
};
let s = parse_struct_info(token_stream);

ProcMacroResult::new(TokenStream::new(generate_op_assign_trait_impl(&op, &s)))
}
Loading

0 comments on commit fdf4098

Please sign in to comment.