iowo/crates/lang/src/ast/mod_tree.rs

use std::{
    cell::RefCell,
    collections::HashMap,
    path::{Path, PathBuf},
    rc::Rc,
    sync::Arc,
};

use rowan::ast::{AstNode, SyntaxNodePtr};

use crate::{lst_parser::syntax_kind::SyntaxKind, Lang, SyntaxNode};

use super::{
    error::Error,
    namespace::Registry,
    nodes::{self, Mod, ModBody, ModName, Root},
    path::{ItemPath, OwnedItemPath},
    source_file::SourceFile,
    Loc,
};

pub struct ModuleTree {
    modules: Rc<HashMap<String, Rc<Module>>>,
}

#[derive(Clone)]
struct ModuleParsingContext {
    files: Rc<RefCell<HashMap<PathBuf, SourceFile>>>,
    errors: Rc<RefCell<Vec<Error>>>,
    reg: Rc<RefCell<Registry>>,
    proj_dir: Rc<PathBuf>,
    cur_path: Vec<String>,
    cur_file: PathBuf,
}

impl ModuleParsingContext {
    fn new(entry_path: &Path) -> Self {
        let proj_dir = entry_path.parent().unwrap().to_owned();
        Self {
            files: Rc::new(RefCell::new(HashMap::new())),
            errors: Rc::new(RefCell::new(Vec::new())),
            reg: Rc::new(RefCell::new(Registry::new())),
            proj_dir: Rc::new(entry_path.to_owned()),
            cur_path: Vec::new(),
            cur_file: entry_path.to_owned(),
        }
    }

    fn entry_path(&self) -> PathBuf {
        self.proj_dir.to_path_buf()
    }

    fn make_mod_file_path(&self, mod_path: &[String], mod_name: &str) -> PathBuf {
        let mut entry_path = self.entry_path();
        for el in mod_path {
            entry_path.push(format!("{el}/"));
        }
        entry_path.push(format!("{mod_name}.owo"));
        entry_path
    }

    fn push_cur_path(&mut self, name_to_add: String) {
        self.cur_path.push(name_to_add);
    }

    fn set_cur_file(&mut self, new: PathBuf) {
        self.cur_file = new;
    }

    fn push_errs(&self, errs: Vec<Error>) {
        self.errors.borrow_mut().extend(errs);
    }

    fn parse_child_modules(
        &self,
        node: &SyntaxNode,
        name: Option<String>,
    ) -> HashMap<String, Rc<Module>> {
        let mut ctx = self.clone();
        if let Some(name) = name {
            ctx.push_cur_path(name)
        }
        node.children()
            .filter_map(Mod::cast)
            .filter_map(|m| match Module::parse_mod(m, self.clone()) {
                Ok(module) => Some((module.name(), module)),
                Err(error) => {
                    self.errors.borrow_mut().push(error);
                    None
                }
            })
            .collect::<HashMap<String, Rc<Module>>>()
    }

    fn dissolve(self) -> (HashMap<PathBuf, SourceFile>, Vec<Error>, Registry) {
        let Self {
            files,
            errors,
            mut reg,
            ..
        } = self;
        (files.take(), errors.take(), (&*reg).take())
    }
}

impl ModuleTree {
    pub fn parse_from_main(
        path: &PathBuf,
        main_file: &SourceFile,
    ) -> (Self, HashMap<PathBuf, SourceFile>, Vec<Error>, Registry) {
        let entry_path = path.parent().unwrap().to_owned();
        let ctx = ModuleParsingContext::new(&entry_path);

        let modules = Rc::new(ctx.parse_child_modules(main_file.tree().syntax(), None));

        let (files, errors, reg) = ctx.dissolve();
        (Self { modules }, files, errors, reg)
    }

    pub fn print_tree(&self, lst: &Root) {
        let name = "main";
        print_tree(&name, self.modules.clone(), 0)
    }
}

#[derive(Debug)]
pub struct Module {
    path: Vec<String>,
    name: String,
    kind: ModuleKind,
    children: Rc<HashMap<String, Rc<Module>>>,
    parent: Option<Rc<Module>>,
    body: SyntaxNodePtr<Lang>,
}

impl Module {
    fn parse_mod(module: Mod, mut ctx: ModuleParsingContext) -> Result<Rc<Self>, Error> {
        let children = module
            .syntax()
            .children()
            // .map(|n| n.kind())
            .collect::<Vec<_>>();

        let r;
        if children.len() == 1 {
            let name = &children[0];
            assert_eq!(name.kind(), SyntaxKind::MODULE_NAME);
            r = Self::parse_file_mod(
                name.text().to_string(),
                ctx.clone(),
                Loc::new(ctx.cur_file.clone(), &module),
            );
        } else if children.len() == 2 {
            let name = &children[0];
            assert_eq!(name.kind(), SyntaxKind::MODULE_NAME);
            let body = &children[1];
            assert_eq!(body.kind(), SyntaxKind::MODULE_BODY);
            r = Ok(Self::parse_inline_mod(module, ctx.clone()));
        } else {
            unreachable!()
        }
        r.map(|module| {
            let rc = Rc::new(module);
            ctx.reg.borrow_mut().insert_mod(rc.path(), rc.clone());
            rc
        })
    }

    fn parse_file_mod(
        name: String,
        mut ctx: ModuleParsingContext,
        decl: Loc<Mod>,
    ) -> Result<Self, Error> {
        ctx.set_cur_file(ctx.make_mod_file_path(&ctx.cur_path, &name));

        let source = match std::fs::read_to_string(&ctx.cur_file) {
            Ok(f) => f,
            Err(e) => return Err(Error::FailedToOpenFileMod(ctx.cur_file, e)),
        };

        let (source_file, file_errors) = SourceFile::parse_from(&ctx.cur_file, source);

        let old_path = ctx.cur_path.clone();
        ctx.push_cur_path(name.clone());
        ctx.push_errs(file_errors);

        let children =
            Rc::new(ctx.parse_child_modules(source_file.tree().syntax(), Some(name.clone())));
        let body = SyntaxNodePtr::new(source_file.tree().syntax());

        ctx.files
            .borrow_mut()
            .insert(ctx.cur_file.clone(), source_file);

        Ok(Self {
            path: old_path,
            name,
            kind: ModuleKind::File {
                declaration: decl,
                file_id: ctx.cur_file.clone(),
            },
            children,
            parent: None,
            body,
        })
    }

    fn parse_inline_mod(module: Mod, mut ctx: ModuleParsingContext) -> Self {
        let mut children = module.syntax().children().collect::<Vec<_>>();

        let body = ModBody::cast(children.pop().unwrap()).unwrap();
        let name = ModName::cast(children.pop().unwrap())
            .unwrap()
            .syntax()
            .text()
            .to_string();

        let old_path = ctx.cur_path.clone();
        ctx.push_cur_path(name.clone());
        let children = Rc::new(ctx.parse_child_modules(body.syntax(), Some(name.clone())));

        Self {
            name,
            kind: ModuleKind::Inline(Loc::new(ctx.cur_file.to_owned(), &module)),
            children,
            parent: None,
            path: old_path,
            body: SyntaxNodePtr::new(body.syntax()),
        }
    }

    pub fn path(&self) -> ItemPath {
        let mut p = self.path.clone();
        p.push(self.name());
        let mut r_p = ItemPath::new();
        for i in p {
            r_p = r_p.push(i);
        }
        r_p
    }

    pub fn name(&self) -> String {
        // self.name.to_node(lst.syntax()).syntax().text().to_string()
        self.name.clone()
    }

    pub fn body(&self, files: &HashMap<PathBuf, SourceFile>) -> SyntaxNode {
        match &self.kind {
            ModuleKind::Inline(l) => {
                let file = files.get(l.file()).unwrap();
                self.body.to_node(file.tree().syntax())
            }
            ModuleKind::File { file_id, .. } => files.get(file_id).unwrap().tree().syntax().clone(),
        }
    }
}

fn print_tree(name: &str, children: Rc<HashMap<String, Rc<Module>>>, level: u32) {
    const INDENT_STR: &str = "    ";

    for _ in 0..level {
        print!("{}", INDENT_STR);
    }
    print!("{name}\n");

    // for (name, module) in children.iter() {
    //     ;
    // }
    children
        .iter()
        .for_each(|(name, m)| print_tree(name, m.children.clone(), level + 1))
}

#[derive(Debug)]
enum ModuleKind {
    Inline(Loc<nodes::Mod>),
    File {
        declaration: Loc<nodes::Mod>,
        file_id: PathBuf,
    },
}