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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

/*
 * Copyright 2022-2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use serde::{Deserialize, Serialize};

/// Describes where in policy source code a node in the CST or expression AST
/// occurs.
#[derive(Serialize, Deserialize, Hash, Debug, Clone, PartialEq, Eq)]
pub struct SourceInfo(pub std::ops::Range<usize>);

impl SourceInfo {
    /// Get the start of range.
    pub fn range_start(&self) -> usize {
        self.0.start
    }

    /// Get the end of range.
    pub fn range_end(&self) -> usize {
        self.0.end
    }
}

/// Metadata for our syntax trees
// Note that these derives are likely to need explicit impls as we develop further
#[derive(Debug, Clone)]
pub struct ASTNode<N> {
    /// Main data represented
    pub node: N,

    /// additional information
    pub info: SourceInfo,
}

impl<N> ASTNode<N> {
    /// Create a new Node from main data
    pub fn new(node: N, left: usize, right: usize) -> Self {
        let info = SourceInfo(left..right);
        ASTNode { node, info }
    }

    /// Create a new Node from main data
    pub fn from_source(node: N, info: SourceInfo) -> Self {
        ASTNode { node, info }
    }

    /// like Option.as_ref()
    pub fn as_ref(&self) -> ASTNode<&N> {
        ASTNode {
            node: &self.node,
            info: self.info.clone(),
        }
    }

    /// map the main data, leaving the SourceInfo alone
    pub fn map<D>(self, f: impl FnOnce(N) -> D) -> ASTNode<D> {
        ASTNode {
            node: f(self.node),
            info: self.info,
        }
    }

    /// consume the Node, producing the main data and the SourceInfo
    pub fn into_inner(self) -> (N, SourceInfo) {
        (self.node, self.info)
    }
}

// Ignore the metadata this node contains
impl<N: PartialEq> PartialEq for ASTNode<N> {
    /// ignores metadata
    fn eq(&self, other: &Self) -> bool {
        self.node == other.node
    }
}
impl<N: Eq> Eq for ASTNode<N> {}

/// Convenience methods on `ASTNode<Option<T>>`
impl<T> ASTNode<Option<T>> {
    /// Similar to `.as_inner()`, but also gives access to the `SourceInfo`
    pub fn as_inner_pair(&self) -> (&SourceInfo, Option<&T>) {
        (&self.info, self.node.as_ref())
    }

    /// Get the inner data as `&T`, if it exists
    pub fn as_inner(&self) -> Option<&T> {
        self.node.as_ref()
    }

    /// `None` if the node is empty, otherwise a node without the `Option`
    pub fn collapse(&self) -> Option<ASTNode<&T>> {
        self.node.as_ref().map(|node| ASTNode {
            node,
            info: self.info.clone(),
        })
    }

    /// Apply the function `f` to the main data and source info. Returns `None`
    /// if no main data or if `f` returns `None`.
    pub fn apply<F, R>(&self, f: F) -> Option<R>
    where
        F: FnOnce(&T, &SourceInfo) -> Option<R>,
    {
        f(self.node.as_ref()?, &self.info)
    }

    /// Apply the function `f` to the main data and source info, consuming them.
    /// Returns `None` if no main data or if `f` returns `None`.
    pub fn into_apply<F, R>(self, f: F) -> Option<R>
    where
        F: FnOnce(T, SourceInfo) -> Option<R>,
    {
        f(self.node?, self.info)
    }
}