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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

// -*- coding: utf-8 -*-
// ------------------------------------------------------------------------------------------------
// Copyright © 2022, tree-sitter authors.
// Licensed under either of Apache License, Version 2.0, or MIT license, at your option.
// Please see the LICENSE-APACHE or LICENSE-MIT files in this distribution for license details.
// ------------------------------------------------------------------------------------------------

use std::collections::hash_map::Entry::Occupied;
use std::collections::hash_map::Entry::Vacant;
use std::collections::HashMap;
use thiserror::Error;

use crate::graph::Value;
use crate::Identifier;

#[derive(Debug, Error)]
pub enum VariableError {
    #[error("Cannot assign immutable variable")]
    CannotAssignImmutableVariable(String),
    #[error("Variable already defined")]
    VariableAlreadyDefined(String),
    #[error("Undefined variable")]
    UndefinedVariable(String),
}

/// An environment of named variables
pub(crate) trait Variables<V> {
    /// Adds a new variable to this environment, returning an error if the variable already
    /// exists.
    fn add(&mut self, name: Identifier, value: V, mutable: bool) -> Result<(), VariableError>;

    /// Sets the variable, returning an error if it does not exists in this environment.
    fn set(&mut self, name: Identifier, value: V) -> Result<(), VariableError>;

    /// Returns the value of a variable, if it exists in this environment.
    fn get(&self, name: &Identifier) -> Option<&V>;
}

/// A map-like implementation of an environment of named variables
pub(crate) struct VariableMap<'a, V> {
    context: Option<&'a mut dyn Variables<V>>,
    values: HashMap<Identifier, Variable<V>>,
}

struct Variable<V> {
    value: V,
    mutable: bool,
}

impl<'a, V> VariableMap<'a, V> {
    /// Creates a new, empty variable environment.
    pub(crate) fn new() -> Self {
        Self {
            context: None,
            values: HashMap::new(),
        }
    }

    /// Creates a nested variable environment, that inherits from the given
    /// context environment.
    pub(crate) fn nested(context: &'a mut dyn Variables<V>) -> Self {
        Self {
            context: Some(context),
            values: HashMap::new(),
        }
    }

    /// Remove a variable from this enviroment, if it exists.
    pub(crate) fn remove(&mut self, name: &Identifier) {
        self.values.remove(name);
    }

    /// Clears this enviroment.
    pub(crate) fn clear(&mut self) {
        self.values.clear();
    }
}

impl<V> Variables<V> for VariableMap<'_, V> {
    fn add(&mut self, name: Identifier, value: V, mutable: bool) -> Result<(), VariableError> {
        match self.values.entry(name) {
            Vacant(v) => {
                let variable = Variable { value, mutable };
                v.insert(variable);
                Ok(())
            }
            Occupied(o) => Err(VariableError::VariableAlreadyDefined(o.key().to_string())),
        }
    }

    fn set(&mut self, name: Identifier, value: V) -> Result<(), VariableError> {
        match self.values.entry(name) {
            Vacant(v) => self
                .context
                .as_mut()
                .map(|context| context.set(v.key().clone(), value))
                .unwrap_or(Err(VariableError::UndefinedVariable(
                    v.into_key().to_string(),
                ))),
            Occupied(mut o) => {
                let mut variable = o.get_mut();
                if variable.mutable {
                    variable.value = value;
                    Ok(())
                } else {
                    Err(VariableError::CannotAssignImmutableVariable(
                        o.key().to_string(),
                    ))
                }
            }
        }
    }

    fn get(&self, name: &Identifier) -> Option<&V> {
        self.values
            .get(name)
            .map(|v| &v.value)
            .or_else(|| self.context.as_ref().map(|p| p.get(name)).flatten())
    }
}

/// Environment of immutable variables
pub struct Globals<'a>(VariableMap<'a, Value>);

impl<'a> Globals<'a> {
    /// Creates a new, empty variable environment.
    pub fn new() -> Self {
        Self(VariableMap::new())
    }

    /// Creates a nested variable environment, that inherits from the given
    /// context environment.
    pub fn nested(context: &'a mut Globals<'a>) -> Self {
        Self(VariableMap::nested(&mut context.0))
    }

    /// Adds a new variable to this environment, returning an error if the variable already
    /// exists.
    pub fn add(&mut self, name: Identifier, value: Value) -> Result<(), VariableError> {
        self.0.add(name, value, false)
    }

    /// Returns the value of a variable, if it exists in this environment.
    pub fn get(&self, name: &Identifier) -> Option<&Value> {
        self.0.get(name)
    }

    /// Remove a variable from this enviroment, if it exists.
    pub fn remove(&mut self, name: &Identifier) {
        self.0.remove(name)
    }

    /// Clears this enviroment.
    pub fn clear(&mut self) {
        self.0.clear()
    }
}