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

use std::collections::HashMap;
use std::fmt::Debug;

use anyhow::Result;
use prql_ast::{expr::Ident, stmt::QueryDef, Span};
use serde::{Deserialize, Serialize};

use super::decl::DeclKind;
use super::{
    decl::{Decl, TableExpr},
    Module, NS_MAIN, NS_QUERY_DEF,
};
use super::{NS_PARAM, NS_STD, NS_THAT, NS_THIS};

/// Context of the pipeline.
#[derive(Serialize, Deserialize, Clone)]
pub struct RootModule {
    /// Map of all accessible names (for each namespace)
    pub(crate) root_mod: Module,

    pub(crate) span_map: HashMap<usize, Span>,
}

// TODO: impl Deref<Target=Module> for RootModule and DerefMut

impl RootModule {
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        // Each module starts with a default namespace that contains a wildcard
        // and the standard library.
        RootModule {
            root_mod: Module {
                names: HashMap::from([
                    (
                        "default_db".to_string(),
                        Decl::from(DeclKind::Module(Module::new_database())),
                    ),
                    (NS_STD.to_string(), Decl::from(DeclKind::default())),
                ]),
                shadowed: None,
                redirects: vec![
                    Ident::from_name(NS_THIS),
                    Ident::from_name(NS_THAT),
                    Ident::from_name(NS_PARAM),
                    Ident::from_name(NS_STD),
                ],
            },
            span_map: HashMap::new(),
        }
    }

    /// Finds that main pipeline given a path to either main itself or its parent module.
    /// Returns main expr and fq ident of the decl.
    pub fn find_main_rel(&self, path: &[String]) -> Result<(&TableExpr, Ident), Option<String>> {
        let (decl, ident) = self.find_main(path)?;

        let decl = (decl.kind.as_table_decl())
            .ok_or(Some(format!("{ident} is not a relational variable")))?;

        Ok((&decl.expr, ident))
    }

    pub fn find_main(&self, path: &[String]) -> Result<(&Decl, Ident), Option<String>> {
        let mut tried_idents = Vec::new();

        // is path referencing the relational var directly?
        if !path.is_empty() {
            let ident = Ident::from_path(path.to_vec());
            let decl = self.root_mod.get(&ident);

            if let Some(decl) = decl {
                return Ok((decl, ident));
            } else {
                tried_idents.push(ident.to_string());
            }
        }

        // is path referencing the parent module?
        {
            let mut path = path.to_vec();
            path.push(NS_MAIN.to_string());

            let ident = Ident::from_path(path);
            let decl = self.root_mod.get(&ident);

            if let Some(decl) = decl {
                return Ok((decl, ident));
            } else {
                tried_idents.push(ident.to_string());
            }
        }

        Err(Some(format!(
            "Expected a declaration at {}",
            tried_idents.join(" or ")
        )))
    }

    pub fn find_query_def(&self, main: &Ident) -> Option<&QueryDef> {
        let ident = Ident {
            path: main.path.clone(),
            name: NS_QUERY_DEF.to_string(),
        };

        let decl = self.root_mod.get(&ident)?;
        decl.kind.as_query_def()
    }

    /// Finds all main pipelines.
    pub fn find_mains(&self) -> Vec<Ident> {
        self.root_mod.find_by_suffix(NS_MAIN)
    }
}

impl Debug for RootModule {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.root_mod.fmt(f)
    }
}