Code from github.com/anakrish/rego-rs

Authored by anakrish and mingweishih

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>

src/ast.rs

@@ -0,0 +1,284 @@
+// Copyright (c) Rego-Rs Authors.
+// Licensed under the Apache 2.0 license.
+
+use crate::lexer::*;
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub enum BinOp {
+    And,
+    Or,
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub enum ArithOp {
+    Add,
+    Sub,
+    Mul,
+    Div,
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub enum BoolOp {
+    Lt,
+    Le,
+    Eq,
+    Ge,
+    Gt,
+    Ne,
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub enum AssignOp {
+    Eq,
+    ColEq,
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub enum Expr<'source> {
+    // Simple items that only have a span as content.
+    String(Span<'source>),
+    RawString(Span<'source>),
+    Number(Span<'source>),
+    True(Span<'source>),
+    False(Span<'source>),
+    Null(Span<'source>),
+    Var(Span<'source>),
+
+    // array
+    Array {
+        span: Span<'source>,
+        items: Vec<Expr<'source>>,
+    },
+
+    // set
+    Set {
+        span: Span<'source>,
+        items: Vec<Expr<'source>>,
+    },
+
+    Object {
+        span: Span<'source>,
+        fields: Vec<(Span<'source>, Expr<'source>, Expr<'source>)>,
+    },
+
+    // Comprehensions
+    ArrayCompr {
+        span: Span<'source>,
+        term: Box<Expr<'source>>,
+        query: Query<'source>,
+    },
+
+    SetCompr {
+        span: Span<'source>,
+        term: Box<Expr<'source>>,
+        query: Query<'source>,
+    },
+
+    ObjectCompr {
+        span: Span<'source>,
+        key: Box<Expr<'source>>,
+        value: Box<Expr<'source>>,
+        query: Query<'source>,
+    },
+
+    Call {
+        span: Span<'source>,
+        fcn: Box<Expr<'source>>,
+        params: Vec<Expr<'source>>,
+    },
+
+    UnaryExpr {
+        span: Span<'source>,
+        expr: Box<Expr<'source>>,
+    },
+
+    // ref
+    RefDot {
+        span: Span<'source>,
+        refr: Box<Expr<'source>>,
+        field: Span<'source>,
+    },
+
+    RefBrack {
+        span: Span<'source>,
+        refr: Box<Expr<'source>>,
+        index: Box<Expr<'source>>,
+    },
+
+    // Infix expressions
+    BinExpr {
+        span: Span<'source>,
+        op: BinOp,
+        lhs: Box<Expr<'source>>,
+        rhs: Box<Expr<'source>>,
+    },
+    BoolExpr {
+        span: Span<'source>,
+        op: BoolOp,
+        lhs: Box<Expr<'source>>,
+        rhs: Box<Expr<'source>>,
+    },
+
+    ArithExpr {
+        span: Span<'source>,
+        op: ArithOp,
+        lhs: Box<Expr<'source>>,
+        rhs: Box<Expr<'source>>,
+    },
+
+    AssignExpr {
+        span: Span<'source>,
+        op: AssignOp,
+        lhs: Box<Expr<'source>>,
+        rhs: Box<Expr<'source>>,
+    },
+
+    Membership {
+        span: Span<'source>,
+        key: Box<Expr<'source>>,
+        value: Box<Option<Expr<'source>>>,
+        collection: Box<Expr<'source>>,
+    },
+}
+
+impl<'source> Expr<'source> {
+    pub fn span(&self) -> &Span<'source> {
+        use Expr::*;
+        match self {
+            String(s) | RawString(s) | Number(s) | True(s) | False(s) | Null(s) | Var(s) => s,
+            Array { span, .. }
+            | Set { span, .. }
+            | Object { span, .. }
+            | ArrayCompr { span, .. }
+            | SetCompr { span, .. }
+            | ObjectCompr { span, .. }
+            | Call { span, .. }
+            | UnaryExpr { span, .. }
+            | RefDot { span, .. }
+            | RefBrack { span, .. }
+            | BinExpr { span, .. }
+            | BoolExpr { span, .. }
+            | ArithExpr { span, .. }
+            | AssignExpr { span, .. }
+            | Membership { span, .. } => span,
+        }
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub enum Literal<'source> {
+    SomeVars {
+        span: Span<'source>,
+        vars: Vec<Span<'source>>,
+    },
+    SomeIn {
+        span: Span<'source>,
+        key: Expr<'source>,
+        value: Option<Expr<'source>>,
+        collection: Expr<'source>,
+    },
+    Expr {
+        span: Span<'source>,
+        expr: Expr<'source>,
+    },
+    NotExpr {
+        span: Span<'source>,
+        expr: Expr<'source>,
+    },
+    Every {
+        span: Span<'source>,
+        key: Span<'source>,
+        value: Option<Span<'source>>,
+        domain: Expr<'source>,
+        query: Query<'source>,
+    },
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub struct WithModifier<'source> {
+    pub span: Span<'source>,
+    pub refr: Expr<'source>,
+    pub r#as: Expr<'source>,
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub struct LiteralStmt<'source> {
+    pub span: Span<'source>,
+    pub literal: Literal<'source>,
+    pub with_mods: Vec<WithModifier<'source>>,
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub struct Query<'source> {
+    pub span: Span<'source>,
+    pub stmts: Vec<LiteralStmt<'source>>,
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub struct RuleAssign<'source> {
+    pub span: Span<'source>,
+    pub op: AssignOp,
+    pub value: Expr<'source>,
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub struct RuleBody<'source> {
+    pub span: Span<'source>,
+    pub assign: Option<RuleAssign<'source>>,
+    pub query: Query<'source>,
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub enum RuleHead<'source> {
+    Compr {
+        span: Span<'source>,
+        refr: Expr<'source>,
+        assign: Option<RuleAssign<'source>>,
+    },
+    Set {
+        span: Span<'source>,
+        refr: Expr<'source>,
+        key: Option<Expr<'source>>,
+    },
+    Func {
+        span: Span<'source>,
+        refr: Expr<'source>,
+        args: Vec<Expr<'source>>,
+        assign: Option<RuleAssign<'source>>,
+    },
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub enum Rule<'source> {
+    Spec {
+        span: Span<'source>,
+        head: RuleHead<'source>,
+        bodies: Vec<RuleBody<'source>>,
+    },
+    Default {
+        span: Span<'source>,
+        refr: Expr<'source>,
+        op: AssignOp,
+        value: Expr<'source>,
+    },
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub struct Package<'source> {
+    pub span: Span<'source>,
+    pub refr: Expr<'source>,
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub struct Import<'source> {
+    pub span: Span<'source>,
+    pub refr: Expr<'source>,
+    pub r#as: Option<Span<'source>>,
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub struct Module<'source> {
+    pub package: Package<'source>,
+    pub imports: Vec<Import<'source>>,
+    pub policy: Vec<Rule<'source>>,
+}

src/builtins/compare.rs

@@ -0,0 +1,27 @@
+use crate::ast::*;
+use crate::value::{Value::*, *};
+
+pub fn eq(v1: &Value, v2: &Value) -> Value {
+    match (v1, v2) {
+        (Undefined, _) | (_, Undefined) => Undefined,
+        _ => Bool(v1 == v2),
+    }
+}
+
+pub fn ne(v1: &Value, v2: &Value) -> Value {
+    match (v1, v2) {
+        (Undefined, _) | (_, Undefined) => Undefined,
+        _ => Bool(v1 != v2),
+    }
+}
+
+pub fn compare(op: &BoolOp, v1: &Value, v2: &Value) -> Value {
+    match op {
+        BoolOp::Eq => eq(v1, v2),
+        BoolOp::Ne => ne(v1, v2),
+        BoolOp::Ge => Bool(v1 >= v2),
+        BoolOp::Gt => Bool(v1 > v2),
+        BoolOp::Le => Bool(v1 <= v2),
+        BoolOp::Lt => Bool(v1 < v2),
+    }
+}

src/builtins/mod.rs

@@ -0,0 +1,3 @@
+mod compare;
+
+pub use self::compare::*;

src/lexer.rs

@@ -0,0 +1,506 @@
+// Copyright (c) Rego-Rs Authors.
+// Licensed under the Apache 2.0 license.
+
+use core::fmt::{Debug, Formatter};
+use core::iter::Peekable;
+use core::str::CharIndices;
+
+use crate::value::Value;
+use anyhow::{anyhow, bail, Result};
+
+#[derive(PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub struct Source<'source> {
+    pub file: &'source str,
+    pub contents: &'source str,
+    pub lines: Vec<&'source str>,
+}
+
+impl<'source> Source<'source> {
+    pub fn message(&self, line: u16, col: u16, kind: &str, msg: &str) -> String {
+        if line as usize > self.lines.len() {
+            return format!("{}: invalid line {} specified", self.file, line);
+        }
+
+        let line_str = format!("{}", line);
+        let line_num_width = line_str.len() + 1;
+        let col_spaces = col as usize - 1;
+
+        format!(
+            "\n-->{}:{}:{}\n{:<line_num_width$}|\n\
+		{:<line_num_width$}| {}\n\
+		{:<line_num_width$}| {:<col_spaces$}^\n\
+		{}: {}",
+            self.file,
+            line,
+            col,
+            "",
+            line,
+            self.lines[line as usize - 1],
+            "",
+            "",
+            kind,
+            msg
+        )
+    }
+
+    pub fn error(&self, line: u16, col: u16, msg: &str) -> anyhow::Error {
+        anyhow!(self.message(line, col, "error", msg))
+    }
+}
+
+#[derive(PartialEq, Eq, PartialOrd, Ord, Clone)]
+pub struct Span<'source> {
+    pub source: &'source Source<'source>,
+    pub line: u16,
+    pub col: u16,
+    pub start: u16,
+    pub end: u16,
+}
+
+impl<'source> Span<'source> {
+    pub fn text(&self) -> &'source str {
+        &self.source.contents[self.start as usize..self.end as usize]
+    }
+}
+
+impl<'source> Debug for Span<'source> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
+        let t = self.text().escape_debug().to_string();
+        let max = 32;
+        let (txt, trailer) = if t.len() > max {
+            (&t[0..max], "...")
+        } else {
+            (t.as_str(), "")
+        };
+
+        f.write_fmt(format_args!(
+            "{}:{}:{}:{}, \"{}{}\"",
+            self.line, self.col, self.start, self.end, txt, trailer
+        ))
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub enum TokenKind {
+    Symbol,
+    String,
+    RawString,
+    Number,
+    Ident,
+    Eof,
+}
+
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub struct Token<'source>(pub TokenKind, pub Span<'source>);
+
+#[derive(Clone)]
+pub struct Lexer<'source> {
+    source: &'source Source<'source>,
+    iter: Peekable<CharIndices<'source>>,
+    line: u16,
+    col: u16,
+}
+
+impl<'source> Lexer<'source> {
+    pub fn new(source: &'source Source<'source>) -> Self {
+        Self {
+            source,
+            iter: source.contents.char_indices().peekable(),
+            line: 1,
+            col: 1,
+        }
+    }
+
+    fn peek(&mut self) -> (usize, char) {
+        match self.iter.peek() {
+            Some((index, chr)) => (*index, *chr),
+            _ => (self.source.contents.len(), '\x00'),
+        }
+    }
+
+    fn peekahead(&mut self, n: usize) -> (usize, char) {
+        match self.iter.clone().nth(n) {
+            Some((index, chr)) => (index, chr),
+            _ => (self.source.contents.len(), '\x00'),
+        }
+    }
+
+    fn read_ident(&mut self) -> Result<Token<'source>> {
+        let start = self.peek().0;
+        let col = self.col;
+        loop {
+            let ch = self.peek().1;
+            if ch.is_ascii_alphanumeric() || ch == '_' {
+                self.iter.next();
+            } else {
+                break;
+            }
+        }
+        let end = self.peek().0;
+        self.col += (end - start) as u16;
+        Ok(Token(
+            TokenKind::Ident,
+            Span {
+                source: self.source,
+                line: self.line,
+                col,
+                start: start as u16,
+                end: end as u16,
+            },
+        ))
+    }
+
+    fn read_digits(&mut self) {
+        while self.peek().1.is_ascii_digit() {
+            self.iter.next();
+        }
+    }
+
+    // See https://www.json.org/json-en.html for number's grammar
+    fn read_number(&mut self) -> Result<Token<'source>> {
+        let (start, chr) = self.peek();
+        let col = self.col;
+        self.iter.next();
+
+        // Read integer part.
+        if chr != '0' {
+            // Starts with 1.. or 9. Read digits.
+            self.read_digits();
+        }
+
+        // Read fraction part
+        // . must be followed by at least 1 digit.
+        if self.peek().1 == '.' && self.peekahead(1).1.is_ascii_digit() {
+            self.iter.next(); // .
+            self.read_digits();
+        }
+
+        // Read exponent part
+        let ch = self.peek().1;
+        if ch == 'e' || ch == 'E' {
+            self.iter.next();
+            // e must be followed by an optional sign and digits
+            if matches!(self.peek().1, '+' | '-') {
+                self.iter.next();
+            }
+            // Read digits. Absence of digit will be validated by serde later.
+            self.read_digits();
+        }
+
+        let end = self.peek().0;
+        self.col += (end - start) as u16;
+
+        // Check for invalid number.Valid number cannot be followed by
+        // these characters:
+        let ch = self.peek().1;
+        if ch == '_' || ch == '.' || ch.is_ascii_alphanumeric() {
+            return Err(self.source.error(self.line, self.col, "invalid number"));
+        }
+
+        // Ensure that the number is parsable in Rust.
+        match serde_json::from_str::<'source, Value>(&self.source.contents[start..end]) {
+            Ok(_) => (),
+            Err(e) => {
+                let serde_msg = &e.to_string();
+                let msg = match &serde_msg {
+                    m if m.contains("out of range") => "out of range",
+                    m if m.contains("invalid number") => "invalid number",
+                    m if m.contains("expected value") => "expected value",
+                    m if m.contains("trailing characters") => "trailing characters",
+                    m => m.to_owned(),
+                };
+
+                bail!(
+                    "{} {}",
+                    self.source.error(
+                        self.line,
+                        col,
+                        "invalid number. serde_json cannot parse number:"
+                    ),
+                    msg
+                )
+            }
+        }
+
+        Ok(Token(
+            TokenKind::Number,
+            Span {
+                source: self.source,
+                line: self.line,
+                col,
+                start: start as u16,
+                end: end as u16,
+            },
+        ))
+    }
+
+    fn read_raw_string(&mut self) -> Result<Token<'source>> {
+        self.iter.next();
+        self.col += 1;
+        let (start, _) = self.peek();
+        let (line, col) = (self.line, self.col);
+        loop {
+            let (_, ch) = self.peek();
+            self.iter.next();
+            match ch {
+                '`' => {
+                    self.col += 1;
+                    break;
+                }
+                '\x00' => {
+                    return Err(self.source.error(line, col, "unmatched `"));
+                }
+                '\t' => self.col += 4,
+                '\n' => {
+                    self.line += 1;
+                    self.col = 1;
+                }
+                _ => self.col += 1,
+            }
+        }
+        let end = self.peek().0;
+        Ok(Token(
+            TokenKind::RawString,
+            Span {
+                source: self.source,
+                line,
+                col,
+                start: start as u16,
+                end: end as u16 - 1,
+            },
+        ))
+    }
+
+    fn read_string(&mut self) -> Result<Token<'source>> {
+        let (line, col) = (self.line, self.col);
+        self.iter.next();
+        self.col += 1;
+        let (start, _) = self.peek();
+        loop {
+            let (offset, ch) = self.peek();
+            let col = self.col + (offset - start) as u16;
+            match ch {
+                '"' | '#' | '\x00' => {
+                    break;
+                }
+                '\\' => {
+                    self.iter.next();
+                    let (_, ch) = self.peek();
+                    self.iter.next();
+                    match ch {
+                        // json escape sequence
+                        '"' | '\\' | '/' | 'b' | 'f' | 'n' | 'r' | 't' => (),
+                        'u' => {
+                            for _i in 0..4 {
+                                let (offset, ch) = self.peek();
+                                let col = self.col + (offset - start) as u16;
+                                if !ch.is_ascii_hexdigit() {
+                                    return Err(self.source.error(
+                                        line,
+                                        col,
+                                        "invalid hex escape sequence",
+                                    ));
+                                }
+                                self.iter.next();
+                            }
+                        }
+                        _ => return Err(self.source.error(line, col, "invalid escape sequence")),
+                    }
+                }
+                _ => {
+                    // check for valid json chars
+                    let col = self.col + (offset - start) as u16;
+                    if !('\u{0020}'..='\u{10FFFF}').contains(&ch) {
+                        return Err(self.source.error(line, col, "invalid character in string"));
+                    }
+                    self.iter.next();
+                }
+            }
+        }
+
+        if self.peek().1 != '"' {
+            return Err(self.source.error(line, col, "unmatched \""));
+        }
+
+        self.iter.next();
+        let end = self.peek().0;
+        self.col += (end - start) as u16;
+
+        // Ensure that the string is parsable in Rust.
+        match serde_json::from_str::<'source, String>(&self.source.contents[start - 1..end]) {
+            Ok(_) => (),
+            Err(e) => {
+                let serde_msg = &e.to_string();
+                let msg = serde_msg;
+                bail!(
+                    "{} {}",
+                    self.source
+                        .error(self.line, col, "serde_json cannot parse string:"),
+                    msg
+                )
+            }
+        }
+
+        Ok(Token(
+            TokenKind::String,
+            Span {
+                source: self.source,
+                line,
+                col: col + 1,
+                start: start as u16,
+                end: end as u16 - 1,
+            },
+        ))
+    }
+
+    fn skip_ws(&mut self) -> Result<()> {
+        // Only the 4 json whitespace characters are recognized.
+        // https://www.crockford.com/mckeeman.html.
+        // Additionally, comments are also skipped.
+        // A tab is considered 4 space characters.
+        'outer: loop {
+            match self.peek().1 {
+                ' ' => self.col += 1,
+                '\t' => self.col += 4,
+                '\r' => {
+                    if self.peekahead(1).1 != '\n' {
+                        return Err(self.source.error(
+                            self.line,
+                            self.col,
+                            "\\r must be followed by \\n",
+                        ));
+                    }
+                }
+                '\n' => {
+                    self.col = 1;
+                    self.line += 1;
+                }
+                '#' => {
+                    self.iter.next();
+                    loop {
+                        match self.peek().1 {
+                            '\n' | '\x00' => continue 'outer,
+                            _ => self.iter.next(),
+                        };
+                    }
+                }
+                _ => break,
+            }
+            self.iter.next();
+        }
+        Ok(())
+    }
+
+    pub fn next_token(&mut self) -> Result<Token<'source>> {
+        self.skip_ws()?;
+
+        let (start, chr) = self.peek();
+        let col = self.col;
+
+        match chr {
+	    // Special case for - followed by digit which is a
+	    // negative json number.
+	    // . followed by digit is invalid number.
+	    '-' | '.' if self.peekahead(1).1.is_ascii_digit() => {
+		self.read_number()
+	    }
+	    // grouping characters
+	    '{' | '}' | '[' | ']' | '(' | ')' |
+	    // arith operator
+	    '+' | '-' | '*' | '/' |
+	    // bin operator
+	    '&' | '|' |
+	    // separators
+	    ',' | ';' | '.' => {
+		self.col += 1;
+		self.iter.next();
+		Ok(Token(TokenKind::Symbol, Span {
+		    source: self.source,
+		    line: self.line,
+		    col,
+		    start: start as u16,
+		    end: start as u16 + 1,
+		}))
+	    }
+	    ':' => {
+		self.col += 1;
+		self.iter.next();
+		let mut end = start as u16 + 1;
+		if self.peek().1 == '=' {
+		    self.col += 1;
+		    self.iter.next();
+		    end += 1;
+		}
+		Ok(Token(TokenKind::Symbol, Span {
+		    source: self.source,
+		    line: self.line,
+		    col,
+		    start: start as u16,
+		    end
+		}))
+	    }
+	    // < <= > >= = ==
+	    '<' | '>' | '=' => {
+		self.col += 1;
+		self.iter.next();
+		if self.peek().1 == '=' {
+		    self.col += 1;
+		    self.iter.next();
+		};
+		Ok(Token(TokenKind::Symbol, Span {
+		    source: self.source,
+		    line: self.line,
+		    col,
+		    start: start as u16,
+		    end: self.peek().0 as u16,
+		}))
+	    }
+	    '!' if self.peekahead(1).1 == '=' => {
+		self.col += 2;
+		self.iter.next();
+		self.iter.next();
+		Ok(Token(TokenKind::Symbol, Span {
+		    source: self.source,
+		    line: self.line,
+		    col,
+		    start: start as u16,
+		    end: self.peek().0 as u16,
+		}))
+	    }
+	    '"' => self.read_string(),
+	    '`' => self.read_raw_string(),
+	    '\x00' => Ok(Token(TokenKind::Eof, Span {
+		source: self.source,
+		line:self.line,
+		col,
+		start: start as u16,
+		end: start as u16
+	    })),
+	    _ if chr.is_ascii_digit() => self.read_number(),
+	    _ if chr.is_ascii_alphabetic() || chr == '_' => {
+		let mut ident = self.read_ident()?;
+		if ident.1.text() == "set" && self.peek().1 == '(' {
+		    // set immediately followed by ( is treated as set( if
+		    // the next token is ).
+		    let state = (self.iter.clone(), self.line, self.col);
+		    self.iter.next();
+
+		    // Check it next token is ).
+		    let next_tok = self.next_token()?;
+		    let is_setp = next_tok.1.text() == ")";
+
+		    // Restore state
+		    (self.iter, self.line, self.col) = state;
+
+		    if is_setp {
+			self.iter.next();
+			self.col += 1;
+			ident.1.end += 1;
+		    }
+		}
+		Ok(ident)
+	    }
+	    _ => Err(self.source.error(self.line, self.col, "invalid character"))
+	}
+    }
+}

src/lib.rs

@@ -0,0 +1,15 @@
+// Copyright (c) Rego-Rs Authors.
+// Licensed under the Apache 2.0 license.
+
+pub mod ast;
+pub mod builtins;
+pub mod interpreter;
+pub mod lexer;
+pub mod parser;
+pub mod value;
+
+pub use ast::*;
+pub use interpreter::*;
+pub use lexer::*;
+pub use parser::*;
+pub use value::*;

src/value.rs

@@ -0,0 +1,325 @@
+// Copyright (c) Rego-Rs Authors.
+// Licensed under the Apache 2.0 license.
+
+use core::fmt;
+use std::collections::{BTreeMap, BTreeSet};
+use std::ops;
+use std::rc::Rc;
+
+use anyhow::{anyhow, Result};
+use ordered_float::OrderedFloat;
+use serde::de::{self, Deserializer};
+use serde::ser::{SerializeMap, Serializer};
+use serde::{Deserialize, Serialize};
+
+// TODO: rego uses BigNum which has arbitrary precision. But there seems
+// to be some bugs with it e.g ((a + b) -a) == b doesn't return true for large
+// values of a and b.
+// Json doesn't specify a limit on precision, but in practice double (f64) seems
+// to be enough to support most use cases and portability too.
+// See discussions in jq's repository.
+// For now we use OrderedFloat<f64>. We can't use f64 directly since it doesn't
+// implement Ord trait.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+pub struct Number(pub OrderedFloat<f64>);
+
+impl Serialize for Number {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let n_f64 = self.0 .0;
+        let n_i64 = n_f64 as i64;
+        let n_u64 = n_f64 as u64;
+
+        if n_u64 as f64 == n_f64 {
+            serializer.serialize_u64(n_u64)
+        } else if n_i64 as f64 == n_f64 {
+            serializer.serialize_i64(n_i64)
+        } else {
+            serializer.serialize_f64(n_f64)
+        }
+    }
+}
+
+struct NumberVisitor;
+impl<'de> de::Visitor<'de> for NumberVisitor {
+    type Value = Number;
+
+    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+        write!(formatter, "a json number")
+    }
+
+    fn visit_f64<E>(self, v: f64) -> Result<Self::Value, E> {
+        Ok(Number(OrderedFloat(v)))
+    }
+
+    fn visit_u64<E>(self, v: u64) -> Result<Self::Value, E> {
+        Ok(Number(OrderedFloat(v as f64)))
+    }
+
+    fn visit_i64<E>(self, v: i64) -> Result<Self::Value, E> {
+        Ok(Number(OrderedFloat(v as f64)))
+    }
+}
+
+impl<'de> Deserialize<'de> for Number {
+    fn deserialize<D>(deserializer: D) -> Result<Number, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        deserializer.deserialize_f64(NumberVisitor)
+    }
+}
+
+impl fmt::Display for Number {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+// We cannot use serde_json::Value because Rego has set type and object's key can be
+// other rego values.
+// BTree is more efficient that a hast table. Another alternative is a sorted vector.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Deserialize)]
+#[serde(untagged)]
+pub enum Value {
+    // Json data types. serde will automatically map json to these variants.
+    Null,
+    Bool(bool),
+    String(String),
+    Number(Number),
+    Array(Rc<Vec<Value>>),
+    Object(Rc<BTreeMap<Value, Value>>),
+
+    // Extra rego data type
+    Set(Rc<BTreeSet<Value>>),
+
+    // Indicate that a value is undefined
+    Undefined,
+}
+
+impl Serialize for Value {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        use serde::ser::Error;
+        match self {
+            Value::Null => serializer.serialize_none(),
+            Value::Bool(b) => serializer.serialize_bool(*b),
+            Value::String(s) => serializer.serialize_str(s.as_str()),
+            Value::Number(n) => n.serialize(serializer),
+            Value::Array(a) => a.serialize(serializer),
+            Value::Object(fields) => {
+                let mut map = serializer.serialize_map(Some(fields.len()))?;
+                for (k, v) in fields.iter() {
+                    match k {
+                        Value::String(_) => map.serialize_entry(k, v)?,
+                        _ => {
+                            let key_str = serde_json::to_string(k).map_err(Error::custom)?;
+                            map.serialize_entry(&key_str, v)?
+                        }
+                    }
+                }
+                map.end()
+            }
+
+            // display set as an array
+            Value::Set(s) => s.serialize(serializer),
+
+            // display undefined as a special string
+            Value::Undefined => serializer.serialize_str("<undefined>"),
+        }
+    }
+}
+
+impl fmt::Display for Value {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match serde_json::to_string(self) {
+            Ok(s) => write!(f, "{}", s),
+            Err(_e) => Err(std::fmt::Error),
+        }
+    }
+}
+
+impl Value {
+    pub fn new_object() -> Value {
+        Value::from_map(BTreeMap::new())
+    }
+
+    pub fn new_set() -> Value {
+        Value::from_set(BTreeSet::new())
+    }
+
+    pub fn new_array() -> Value {
+        Value::from_array(vec![])
+    }
+
+    pub fn from_json_str(json: &str) -> Result<Value> {
+        Ok(serde_json::from_str(json)?)
+    }
+
+    pub fn to_json_str(&self) -> Result<String> {
+        Ok(serde_json::to_string_pretty(self)?)
+    }
+}
+
+impl Value {
+    pub fn from_f64(v: f64) -> Value {
+        Value::Number(Number(OrderedFloat(v)))
+    }
+
+    pub fn from_array(a: Vec<Value>) -> Value {
+        Value::Array(Rc::new(a))
+    }
+
+    pub fn from_set(s: BTreeSet<Value>) -> Value {
+        Value::Set(Rc::new(s))
+    }
+
+    pub fn from_map(m: BTreeMap<Value, Value>) -> Value {
+        Value::Object(Rc::new(m))
+    }
+
+    pub fn is_null(&self) -> bool {
+        matches!(self, Value::Null)
+    }
+
+    pub fn is_undefined(&self) -> bool {
+        matches!(self, Value::Null)
+    }
+
+    pub fn as_bool(&self) -> Result<&bool> {
+        match self {
+            Value::Bool(b) => Ok(b),
+            _ => Err(anyhow!("not a bool")),
+        }
+    }
+
+    pub fn as_bool_mut(&mut self) -> Result<&mut bool> {
+        match self {
+            Value::Bool(b) => Ok(b),
+            _ => Err(anyhow!("not a bool")),
+        }
+    }
+
+    pub fn as_string(&self) -> Result<&String> {
+        match self {
+            Value::String(s) => Ok(s),
+            _ => Err(anyhow!("not a string")),
+        }
+    }
+
+    pub fn as_string_mut(&mut self) -> Result<&mut String> {
+        match self {
+            Value::String(s) => Ok(s),
+            _ => Err(anyhow!("not a string")),
+        }
+    }
+
+    pub fn as_number(&self) -> Result<&Number> {
+        match self {
+            Value::Number(n) => Ok(n),
+            _ => Err(anyhow!("not a number")),
+        }
+    }
+
+    pub fn as_number_mut(&mut self) -> Result<&mut Number> {
+        match self {
+            Value::Number(n) => Ok(n),
+            _ => Err(anyhow!("not a number")),
+        }
+    }
+
+    pub fn as_array(&self) -> Result<&Vec<Value>> {
+        match self {
+            Value::Array(a) => Ok(a),
+            _ => Err(anyhow!("not an array")),
+        }
+    }
+
+    pub fn as_array_mut(&mut self) -> Result<&mut Vec<Value>> {
+        match self {
+            Value::Array(a) => Ok(Rc::make_mut(a)),
+            _ => Err(anyhow!("not an array")),
+        }
+    }
+
+    pub fn as_set(&self) -> Result<&BTreeSet<Value>> {
+        match self {
+            Value::Set(s) => Ok(s),
+            _ => Err(anyhow!("not a set")),
+        }
+    }
+
+    pub fn as_set_mut(&mut self) -> Result<&mut BTreeSet<Value>> {
+        match self {
+            Value::Set(s) => Ok(Rc::make_mut(s)),
+            _ => Err(anyhow!("not a set")),
+        }
+    }
+
+    pub fn as_object(&self) -> Result<&BTreeMap<Value, Value>> {
+        match self {
+            Value::Object(m) => Ok(m),
+            _ => Err(anyhow!("not an object")),
+        }
+    }
+
+    pub fn as_object_mut(&mut self) -> Result<&mut BTreeMap<Value, Value>> {
+        match self {
+            Value::Object(m) => Ok(Rc::make_mut(m)),
+            _ => Err(anyhow!("not an object")),
+        }
+    }
+}
+
+impl ops::Index<usize> for Value {
+    type Output = Value;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        match self.as_array() {
+            Ok(a) if index < a.len() => &a[index],
+            _ => &Value::Undefined,
+        }
+    }
+}
+
+impl ops::Index<&str> for Value {
+    type Output = Value;
+
+    fn index(&self, key: &str) -> &Self::Output {
+        &self[&Value::String(key.to_owned())]
+    }
+}
+
+impl ops::Index<&String> for Value {
+    type Output = Value;
+
+    fn index(&self, key: &String) -> &Self::Output {
+        &self[&Value::String(key.clone())]
+    }
+}
+
+impl ops::Index<&Value> for Value {
+    type Output = Value;
+
+    fn index(&self, key: &Value) -> &Self::Output {
+        match (self, &key) {
+            (Value::Object(o), _) => match &o.get(key) {
+                Some(v) => v,
+                _ => &Value::Undefined,
+            },
+            (Value::Array(a), Value::Number(n)) => {
+                let index = n.0 .0 as usize;
+                if index < a.len() {
+                    &a[index]
+                } else {
+                    &Value::Undefined
+                }
+            }
+            _ => &Value::Undefined,
+        }
+    }
+}