deflake-rs 0.1.0

use core::panic;
use std::collections::HashSet;

use duplicate::duplicate_item;
use proc_macro2::Span;
use quote::quote;
use syn::{
    spanned::Spanned,
    visit::{self, Visit},
    File,
};

use crate::git::{Edit, EditType};

pub trait SpanHelpers {
    fn has_line(&self, line: usize) -> bool;
}
impl SpanHelpers for Span {
    fn has_line(&self, line: usize) -> bool {
        let line = line;
        self.start().line <= line && self.end().line >= line
    }
}

pub trait ItemFnHelpers {
    fn attr_lines(&self) -> HashSet<usize>;
    fn trackable_lines(&self) -> HashSet<usize>;
}

#[duplicate_item(
    func_type;
    [syn::ItemFn];
    [syn::ImplItemFn]
)]
impl ItemFnHelpers for func_type {
    fn attr_lines(&self) -> HashSet<usize> {
        let mut lines = HashSet::new();

        // Function attributes
        for attr in &self.attrs {
            let span = attr.span();
            lines.extend(span.start().line..=span.end().line);
        }

        lines
    }

    fn trackable_lines(&self) -> HashSet<usize> {
        let mut lines = HashSet::new();

        // Statement lines
        for stmt in &self.block.stmts {
            let span = stmt.span();
            lines.extend(span.start().line..=span.end().line);
        }

        lines.extend(self.attr_lines());

        // Function signature and braces
        lines.extend(self.sig.span().start().line..=self.sig.span().end().line);

        let braces = self.block.brace_token.span.span();
        lines.insert(braces.start().line);
        lines.insert(braces.end().line);

        lines
    }
}

pub struct AST {}

impl AST {
    pub fn line_in_span(span: Span, line: u32) -> bool {
        let line = line as usize;

        return span.start().line <= line && span.end().line >= line;
    }

    pub fn parse_file(file: &str) -> File {
        let syntax = syn::parse_str::<File>(file).unwrap();

        syntax
    }
}

pub struct AddedVisitor {
    pub functions: Vec<ItemFn>,
    //pub track: Vec<Track>,
}
impl AddedVisitor {
    pub fn new() -> AddedVisitor {
        AddedVisitor { functions: vec![] }
    }
}
impl<'ast> Visit<'ast> for AddedVisitor {
    fn visit_item_fn(&mut self, i: &'ast syn::ItemFn) {
        let span = i.span();
        let fn_name = i.sig.ident.to_string();

        let sig = &i.sig;
        let sig = quote!(#sig).to_string();

        let func = ItemFn {
            name: fn_name,
            span,
            sig,
            ident_line: i.sig.ident.span().start().line,
            trackable_lines: i.trackable_lines(),
            attr_lines: i.attr_lines(),
        };

        self.functions.push(func.clone());
    }

    fn visit_impl_item_fn(&mut self, i: &'ast syn::ImplItemFn) {
        let span = i.span();
        let fn_name = i.sig.ident.to_string();

        let sig = &i.sig;
        let sig = quote!(#sig).to_string();

        let func = ItemFn {
            name: fn_name,
            span,
            sig,
            ident_line: i.sig.ident.span().start().line,
            trackable_lines: i.trackable_lines(),
            attr_lines: i.attr_lines(),
        };

        self.functions.push(func.clone());
    }
}

// Could be anything Rust AST wise, deal with functions, impl blocks and such
#[derive(Debug, Clone)]
pub struct ItemFn {
    pub name: String,
    pub span: Span,
    pub sig: String,
    pub ident_line: usize,
    pub trackable_lines: HashSet<usize>, // lists all the trackable lines in the function
    pub attr_lines: HashSet<usize>,
}

#[derive(Debug, Clone)]
enum Item {
    Fn(ItemFn, Span),
}

impl Item {
    fn difference(a: &Vec<Item>, b: &Vec<Item>) -> Vec<Item> {
        let mut c = vec![];
        for item in a {
            let other_item = (b).into_iter().find(|i| i.eq(&item));

            if None == other_item {
                c.push(item.clone());
            }
        }

        c
    }

    fn functions(v: &Vec<Item>) -> Vec<ItemFn> {
        let mut funcs = vec![];
        for item in v {
            match item {
                Item::Fn(f, _) => funcs.push(f.clone()),
            }
        }

        return funcs;
    }
}

impl PartialEq for ItemFn {
    fn eq(&self, other: &Self) -> bool {
        // TODO: add more fields
        self.name == other.name
    }
}

impl PartialEq for Item {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Item::Fn(x, _), Item::Fn(y, _)) => x.eq(y),
            _ => false,
        }
    }
}

#[derive(Debug, Clone)]
struct Module {
    name: String,
    span: Option<Span>,

    items: Vec<Item>,

    sub_modules: Vec<NestedItem>,
    //stack: Stack,
}

impl Module {
    fn new(name: impl ToString, span: Option<Span>) -> Module {
        Module {
            name: name.to_string(),
            span,
            items: vec![],
            sub_modules: vec![],
            //stack: Stack(vec![]),
        }
    }

    fn difference(a: &Self, b: &Self) -> (Vec<Item>, Vec<NestedItem>) {
        (
            Item::difference(&a.items, &b.items),
            NestedItem::difference(&a.sub_modules, &b.sub_modules),
        )
    }

    fn functions(&self) -> Vec<ItemFn> {
        let mut funcs = vec![];

        funcs.append(&mut Item::functions(&self.items));
        funcs.append(&mut NestedItem::functions(&self.sub_modules));

        funcs
    }
}

impl PartialEq for Module {
    fn eq(&self, other: &Self) -> bool {
        self.name == other.name
            && vec_eq(&self.items, &other.items)
            && vec_eq(&self.sub_modules, &other.sub_modules)
    }
}

#[derive(Debug, Clone)]
struct Impl {
    name: String,
    span: Span,
    items: Vec<Item>,
}

impl Impl {
    fn new(name: impl ToString, span: Span) -> Impl {
        Impl {
            name: name.to_string(),
            span,
            items: vec![],
        }
    }

    fn difference(a: &Self, b: &Self) -> Vec<Item> {
        let mut c = vec![];
        for item in &a.items {
            let other_item = (&b.items).into_iter().find(|i| i.eq(&item));

            if None == other_item {
                c.push(item.clone());
            }
        }

        c
    }

    fn functions(&self) -> Vec<ItemFn> {
        Item::functions(&self.items)
    }
}

impl PartialEq for Impl {
    fn eq(&self, other: &Self) -> bool {
        self.name == other.name && vec_eq(&self.items, &other.items)
    }
}

// TODO: store more info about each module / impl block
#[derive(Debug, Clone)]
pub enum NestedItem {
    Module(Module), // span, mod name, attribute list
    Impl(Impl),     // span, struct name, trait?
}

impl NestedItem {
    fn items(&mut self) -> &mut Vec<Item> {
        match self {
            NestedItem::Impl(i) => &mut i.items,
            NestedItem::Module(m) => &mut m.items,
        }
    }

    fn sub_modules(&mut self) -> &mut Vec<NestedItem> {
        match self {
            NestedItem::Module(m) => &mut m.sub_modules,
            NestedItem::Impl(i) => panic!("Unexpected"),
        }
    }

    fn shallow_eq(&self, other: &Self) -> bool {
        match (self, other) {
            (NestedItem::Module(a), NestedItem::Module(b)) => a.name == b.name,
            (NestedItem::Impl(a), NestedItem::Impl(b)) => a.name == b.name,
            _ => false,
        }
    }

    // Find everything in a but not in b
    fn difference(a: &Vec<Self>, b: &Vec<Self>) -> Vec<Self> {
        let mut c = vec![];
        for item in a {
            // Find matching nested item (based on name)
            let other_item = b.into_iter().find(|m| m.shallow_eq(&item));
            if let Some(other_item) = other_item {
                match (item, other_item) {
                    (NestedItem::Module(a), NestedItem::Module(b)) => {
                        let mut new_item = a.clone();
                        let (items, nested_items) = Module::difference(a, b);

                        if items.len() > 0 || nested_items.len() > 0 {
                            new_item.items = items;
                            new_item.sub_modules = nested_items;
                            c.push(NestedItem::Module(new_item));
                        }
                    }
                    (NestedItem::Impl(a), NestedItem::Impl(b)) => {
                        let mut new_item = a.clone();

                        let items = Impl::difference(a, b);
                        if items.len() > 0 {
                            new_item.items = items;
                            c.push(NestedItem::Impl(new_item));
                        }
                    }
                    _ => (),
                };

                //item.difference(&other_item)
            } else {
                c.push(item.clone());
            }
        }

        c
    }

    pub fn functions(v: &Vec<NestedItem>) -> Vec<ItemFn> {
        let mut f = vec![];

        for i in v {
            let mut item_funcs = match i {
                NestedItem::Impl(imp) => imp.functions(),
                NestedItem::Module(module) => module.functions(),
            };
            f.append(&mut item_funcs);
        }

        f
    }
}

impl PartialEq for NestedItem {
    // TODO: this is really bad. should use a hash table instead for storing items and sub
    // modules IMO
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (NestedItem::Impl(s), NestedItem::Impl(i)) => s.eq(i),
            (NestedItem::Module(s), NestedItem::Module(i)) => s.eq(i),
            _ => false,
        }
    }
}

#[derive(Debug)]
pub struct ModificationVisitor {
    //pub highlevel_track: Vec<Track>,
    new_version: bool,
    pub functions: Vec<ItemFn>, // used for fast detection of if an edit in a function or not
    modules: Vec<NestedItem>,
}
impl ModificationVisitor {
    pub fn new(is_new: bool) -> ModificationVisitor {
        ModificationVisitor {
            //track: vec![],
            new_version: is_new,
            functions: vec![],
            modules: vec![NestedItem::Module(Module::new("_", None))], // Top level functions grouped in a module
        }
    }

    fn current_module(&mut self) -> &mut NestedItem {
        self.modules.last_mut().unwrap()
    }

    // Get items in a but not b
    pub fn difference(&self, other: &Self) -> Vec<NestedItem> {
        let ab = NestedItem::difference(&self.modules, &other.modules);
        //let ba = NestedItem::difference(&other.modules, &self.modules);
        //dbg!(ab);
        //dbg!(ba);
        ab
    }

    // NOTE: uses assumption that we aren't tracking nested functions, so only ever returns highest
    // level function
    pub fn find_function(&self, line: usize) -> Option<ItemFn> {
        for func in &self.functions {
            if func.span.has_line(line) {
                return Some(func.clone());
            }
        }

        None
    }
}

// TODO: very inefficient. hash map compare should be used instead I think
fn vec_eq<T: PartialEq>(a: &Vec<T>, b: &Vec<T>) -> bool {
    let mut a = a.into_iter();
    let mut b = b.into_iter();
    a.all(|x| b.any(|y| x.eq(y))) && b.all(|x| a.any(|y| x.eq(y)))
}

impl PartialEq for ModificationVisitor {
    fn eq(&self, other: &Self) -> bool {
        vec_eq(&self.modules, &other.modules)
    }
}

impl<'ast> Visit<'ast> for ModificationVisitor {
    // TODO: find a way to combine with visit_impl_item_fn
    fn visit_item_fn(&mut self, i: &'ast syn::ItemFn) {
        let span = i.span();
        let fn_name = i.sig.ident.to_string();

        // TODO: hacky for now
        let sig = &i.sig;
        let sig = quote!(#sig).to_string();

        let func = ItemFn {
            name: fn_name,
            span,
            sig,
            ident_line: i.sig.ident.span().start().line,
            trackable_lines: i.trackable_lines(),
            attr_lines: i.attr_lines(),
        };

        self.functions.push(func.clone());

        self.current_module().items().push(Item::Fn(func, span));
    }

    fn visit_impl_item_fn(&mut self, i: &'ast syn::ImplItemFn) {
        let span = i.span();
        let fn_name = i.sig.ident.to_string();

        // TODO: hacky for now
        let sig = &i.sig;
        let sig = quote!(#sig).to_string();

        let func = ItemFn {
            name: fn_name,
            span,
            sig,
            ident_line: i.sig.ident.span().start().line,
            trackable_lines: i.trackable_lines(),
            attr_lines: i.attr_lines(),
        };

        self.functions.push(func.clone());

        self.current_module().items().push(Item::Fn(func, span));
    }

    fn visit_item_mod(&mut self, i: &'ast syn::ItemMod) {
        let ident = &i.ident;
        let mod_name = quote!(#ident).to_string();

        // TODO: store attirbutes too?
        //i.attrs.iter().map(|a| quote!(#a).to_string()).collect(),
        self.modules.push(NestedItem::Module(Module::new(
            mod_name,
            Some(i.span().clone()),
        )));

        visit::visit_item_mod(self, i);

        // Add the finished module to its parent module
        let finished_mod = self.modules.pop().unwrap();
        self.modules
            .last_mut()
            .unwrap()
            .sub_modules()
            .push(finished_mod);
    }

    fn visit_item_impl(&mut self, i: &'ast syn::ItemImpl) {
        let ty = &i.self_ty;
        let ty_name = quote!(#ty).to_string();

        self.modules
            .push(NestedItem::Impl(Impl::new(ty_name, i.span().clone())));

        visit::visit_item_impl(self, i);

        let finished_mod = self.modules.pop().unwrap();
        self.modules
            .last_mut()
            .unwrap()
            .sub_modules()
            .push(finished_mod);
    }
}

// TODO: add a test suite using a local file to create the AST?