1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

pub mod global;
pub mod local;

pub use crate::ir::quantity::{global::GlobalVariableName, local::RegisterName};
use crate::utility::parsing;
use enum_dispatch::enum_dispatch;
use nom::{branch::alt, combinator::map, IResult};
use serde::{Deserialize, Serialize};
use std::fmt::{self, Display, Formatter};
/// Tag trait for [`Quantity`].
#[enum_dispatch]
trait IsQuantity {}

/// [`Quantity`] represents a variable (global or local) or a constant value
#[enum_dispatch(IsQuantity)]
#[derive(Debug, Eq, PartialEq, Clone, Hash, Deserialize, Serialize)]
pub enum Quantity {
    RegisterName,
    GlobalVariableName,
    NumberLiteral(i64),
}

impl Quantity {
    pub fn unwrap_local(self) -> RegisterName {
        if let Quantity::RegisterName(register_name) = self {
            register_name
        } else {
            panic!("Not local")
        }
    }

    pub fn as_local(&self) -> Option<&RegisterName> {
        if let Quantity::RegisterName(register_name) = self {
            Some(register_name)
        } else {
            None
        }
    }
}

impl Display for Quantity {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Quantity::GlobalVariableName(global) => write!(f, "{global}"),
            Quantity::RegisterName(local) => write!(f, "{local}"),
            Quantity::NumberLiteral(number) => write!(f, "{number}"),
        }
    }
}

/// Parse source code to get a [`Quantity`].
pub fn parse(code: &str) -> IResult<&str, Quantity> {
    alt((
        map(local::parse, Quantity::RegisterName),
        map(global::parse, Quantity::GlobalVariableName),
        map(parsing::integer, Quantity::NumberLiteral),
    ))(code)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse() {
        let result = parse("%foo").unwrap().1;
        assert_eq!(
            result,
            Quantity::RegisterName(RegisterName("foo".to_string()))
        );
        let result = parse("%0").unwrap().1;
        assert_eq!(
            result,
            Quantity::RegisterName(RegisterName("0".to_string()))
        );
        let result = parse("@foo").unwrap().1;
        assert_eq!(
            result,
            Quantity::GlobalVariableName(GlobalVariableName("foo".to_string()))
        );
        let result = parse("123").unwrap().1;
        assert_eq!(result, Quantity::NumberLiteral(123));
    }
}