Pesto handle lib.makeOverridable

pesto/src/alias.rs

@@ -123,7 +123,7 @@ pub fn categorize(data: &Vec<Docs>) -> FnCategories {
             if lambda.isFunctor == Some(true) {
                 // A functor takes self as first argument
                 // Subtract the first argument from the count of applied arguments.
-                match lambda.countApplied.map(|s|s-1) {
+                match lambda.countApplied.map(|s| if s != 0 { s - 1 } else { 0 }) {
                     // Some(0) | None => {
                     Some(0) => {
                         if lambda.isPrimop {

pesto/src/bulk.rs

@@ -3,7 +3,10 @@ use std::{collections::HashMap, path::PathBuf, println, rc::Rc, time::Instant, v
 use crate::{
     alias::{categorize, init_alias_map},
     pasta::{Docs, Files, Pasta},
-    position::{DocComment, DocIndex, FilePosition, NixDocComment},
+    position::{
+        get_overridable_fn, seek_file_position, when_overridable_lambda, DocComment, DocIndex,
+        FilePosition, NixDocComment,
+    },
 };
 
 #[derive(Debug)]
@@ -112,6 +115,8 @@ impl<'a> BulkProcessing for Pasta {
         let file_map = build_file_map(&data);
 
         let mut pos_doc_map: HashMap<&FilePosition, Option<NixDocComment>> = HashMap::new();
+
+        let mut file_idx_map: HashMap<&PathBuf, DocIndex> = HashMap::new();
         for (path, lookups) in file_map.iter() {
             if !path.exists() {
                 println!("file does not exist: {:?} Skipping.", path);
@@ -127,6 +132,7 @@ impl<'a> BulkProcessing for Pasta {
             );
 
             let doc_index = DocIndex::new(path, positions);
+            file_idx_map.insert(path, doc_index.clone());
 
             for lookup in lookups {
                 pos_doc_map.insert(
@@ -135,8 +141,29 @@ impl<'a> BulkProcessing for Pasta {
                 );
             }
         }
+
         let mut filled_docs = fill_docs(&data, &pos_doc_map);
 
+        // Do a second pass for potential lib.makeOverridable wrapped functions.
+        let mut restores: Vec<String> = vec![];
+        for doc in filled_docs.iter_mut() {
+            if let Some(orig_file) = when_overridable_lambda(doc, &file_idx_map) {
+                let orig_lambda = get_overridable_fn(&orig_file);
+                let orig_pos =
+                    orig_lambda.map(|n| seek_file_position(&orig_file, &n.text_range().start()));
+                if let Some(orig_pos) = orig_pos {
+                    if let Some(l) = &mut doc.docs.lambda {
+                        l.position = orig_pos;
+                        restores.push(format!("{:?}", doc.path.join(".")));
+                    }
+                }
+            }
+        }
+        println!(
+            "Restored from makeOverridable with edge-case handling: {:?}",
+            restores
+        );
+
         let categories = categorize(&filled_docs);
         let alias_map = init_alias_map(&data, categories);
 

pesto/src/main.rs

@@ -88,7 +88,7 @@ pub fn main() {
             let matter = &document.meta;
             let content = &document.content;
 
-            let signature = content
+            let _signature = content
                 .as_ref()
                 .map(|c| c.content.as_ref().map(|s| find_type(&s)))
                 .flatten();
@@ -132,9 +132,9 @@ pub fn main() {
 /// Find the content which should be displayed.
 /// The own attribute content is the correct one usually.
 /// Sometimes there is no attribute content.
-/// The we search all the aliases for their attribute content.
+/// Then we search all the aliases for their attribute content.
 /// As a fallback we can display the own lambda content.
-fn find_document_content<'a>(
+pub fn find_document_content<'a>(
     item: &'a Docs,
     all: &'a HashMap<Rc<ValuePath>, Docs>,
 ) -> Option<ContentSource<'a>> {

pesto/src/position.rs

@@ -1,3 +1,4 @@
+use expect_test::Position;
 use rnix::ast::{self};
 use rnix::{match_ast, SyntaxNode};
 use rowan::TextSize;
@@ -12,6 +13,8 @@ use std::rc::Rc;
 use std::{format, fs, path::PathBuf, println};
 
 use crate::comment::get_expr_docs;
+use crate::find_document_content;
+use crate::pasta::{Docs, ValuePath};
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct TextPosition {
@@ -26,7 +29,7 @@ pub struct FilePosition {
     pub column: usize,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct DocIndex<'a> {
     file: &'a PathBuf,
     pos_idx: HashMap<(usize, usize), TextSize>,
@@ -36,16 +39,179 @@ pub struct DocIndex<'a> {
 pub trait DocComment<'a> {
     fn new(file: &'a PathBuf, positions: HashMap<usize, Vec<usize>>) -> Self;
     fn get_docs(&self, line: usize, column: usize) -> Option<NixDocComment>;
+    fn get_node_at_position(&self, position: &'a FilePosition) -> &Option<SyntaxNode>;
 }
 
-fn get_src(path: &PathBuf) -> String {
+pub fn get_src(path: &PathBuf) -> String {
     if let Ok(src) = fs::read_to_string(path) {
         return src;
     }
     println!("could not read file: {}", path.to_str().unwrap());
+
     exit(1);
 }
 
+/// Returns the node path
+pub fn get_call_package_file(node: Option<&SyntaxNode>) -> Option<SyntaxNode> {
+    if let Some(node) = node {
+        match node.kind() {
+            rnix::SyntaxKind::NODE_ATTRPATH_VALUE => {
+                get_call_package_file(node.last_child().as_ref())
+            }
+            rnix::SyntaxKind::NODE_PATH => Some(node.clone()),
+            rnix::SyntaxKind::NODE_APPLY => match node.first_child().map(|n| n.kind()) {
+                Some(rnix::SyntaxKind::NODE_APPLY) => {
+                    get_call_package_file(node.first_child().as_ref())
+                }
+                _ => get_call_package_file(node.last_child().as_ref()),
+            },
+            _ => {
+                println!(
+                    "Unhandled node when trying to unpack callPackage expression: {:?}",
+                    node.kind()
+                );
+                None
+            } // n => get_call_package_file(node.last_child().as_ref()),
+        }
+    } else {
+        None
+    }
+}
+
+/// Goes up the tree to find the parent node that matches the predicate,
+/// checks starting with the current node
+/// Stops when the limit is reached
+fn check_outwards<F>(node: &SyntaxNode, pred: F, limit: usize, curr: usize) -> bool
+where
+    F: Fn(&SyntaxNode) -> bool,
+{
+    if pred(node) {
+        return true;
+    } else if curr < limit {
+        node.parent()
+            .map(|parent| check_outwards(&parent, pred, limit, curr + 1))
+            .unwrap_or(false)
+    } else {
+        return false;
+    }
+}
+
+fn match_attrpath_ident(node: &SyntaxNode, text: &str) -> bool {
+    match node.kind() {
+        rnix::SyntaxKind::NODE_ATTRPATH_VALUE => {
+            let maybe_ident = node
+                .first_child()
+                .and_then(|attrpath| attrpath.first_child());
+            if let Some(ident) = maybe_ident {
+                ident.text().to_string() == text
+            } else {
+                false
+            }
+        }
+        _ => false,
+    }
+}
+
+pub fn when_overridable_lambda(
+    doc: &Docs,
+    file_idx_map: &HashMap<&PathBuf, DocIndex>,
+) -> Option<PathBuf> {
+    let all: HashMap<Rc<ValuePath>, Docs> = HashMap::new();
+    match find_document_content(doc, &all) {
+        Some(content) => {
+            if content.content.is_some() && !content.content.as_ref().unwrap().is_empty() {
+                return None;
+            }
+
+            if content.source.is_none() {
+                // println!("Could not find source position for {:?}", &doc.path);
+                return None;
+            }
+            let position = content.source.unwrap().position;
+
+            if position.is_none() {
+                // println!("Could not find source position for {:?}", &doc.path);
+                return None;
+            }
+
+            let position = position.unwrap();
+            let doc_index = file_idx_map.get(&position.file);
+
+            if let Some(doc_index) = doc_index {
+                let node = doc_index.get_node_at_position(position).as_ref();
+                let is_in_overridable = check_outwards(
+                    node.unwrap(),
+                    |node| match node.kind() {
+                        rnix::SyntaxKind::NODE_ATTRPATH_VALUE => {
+                            let check = match_attrpath_ident(node, "makeOverridable");
+                            check
+                        }
+                        _ => false,
+                    },
+                    5,
+                    0,
+                );
+                if !is_in_overridable {
+                    println!("function is not wrapped in makeOverridable {:?}", &doc.path);
+                    return None;
+                }
+                let pos = &doc.docs.attr.position.as_ref();
+                if pos.is_none() {
+                    println!("Could not find source position for {:?}", &doc.path);
+                    return None;
+                }
+                let pos = pos.unwrap();
+
+                println!("Found function that is wrapped in lib.makeOverridable: {:?}. Trying to infer original lambda. From {:?}", &doc.path, &pos);
+
+                let package_file_idx = file_idx_map.get(&pos.file).unwrap();
+                let node = package_file_idx.get_node_at_position(pos);
+                let package_file = node
+                    .as_ref()
+                    .map(|node| {
+                        let is_node = match_attrpath_ident(node, doc.path.last().unwrap());
+                        if !is_node {
+                            println!("Could not find node for {:?}", doc.path.last().unwrap());
+                            return None;
+                        }
+                        let path = get_call_package_file(Some(node));
+                        let package_file = path.map(|v| v.text().to_string());
+                        package_file
+                    })
+                    .flatten();
+
+                if package_file.is_none() {
+                    println!("Could not find package file for {:?}; While trying to cover lib.makeOverridable case handling", &doc.path);
+                    return None;
+                }
+                let rel_package_file = package_file.unwrap();
+
+                let resolved_path = pos
+                    .file
+                    .parent()
+                    .map(|parent_path| {
+                        parent_path
+                            .join(rel_package_file)
+                            .canonicalize()
+                            .ok()
+                            .map(|p| {
+                                if p.to_str().unwrap().ends_with(".nix") {
+                                    p
+                                } else {
+                                    p.join("default.nix")
+                                }
+                            })
+                    })
+                    .flatten();
+                return resolved_path;
+            } else {
+                None
+            }
+        }
+        _ => None,
+    }
+}
+
 /// Initializes a HashMap for lookup operation between L:C and absolute position.
 /// Returns both
 /// Position HashMap from l:c -> abs
@@ -84,6 +250,150 @@ fn init_pos_idx(
     (res, inverse)
 }
 
+/// Inefficient way to get the node at a position
+/// Reads the whole file and iterates over the AST
+/// Returns a SyntaxNode if found
+pub fn get_overridable_fn(file: &PathBuf) -> Option<SyntaxNode> {
+    println!("get_overridable_fn from {:?}", &file);
+    let src: String = get_src(file);
+    let rc: Rc<String> = Rc::new(src);
+    let ast = rnix::Root::parse(Rc::clone(&rc).as_str()).syntax();
+    get_call_package_lambda_body(&ast)
+        .map(|body| {
+            let maybe_wrapped_lambda = get_apply_make_overridable_body(&body);
+            maybe_wrapped_lambda.map(|l| unpack_lambda(&l))
+        })
+        .flatten()
+        .flatten()
+}
+
+pub fn seek_file_position(path: &PathBuf, text_pos: &TextSize) -> Option<FilePosition> {
+    let file = File::open(path).unwrap();
+    let reader = BufReader::new(file);
+
+    let mut curr_position = 0;
+    for (curr_line, line) in reader.lines().enumerate() {
+        let pos_found = match line {
+            Ok(line) => {
+                let absolute = TextSize::from(u32::try_from(curr_position + line.len()).unwrap());
+                let curr = TextSize::from(u32::try_from(curr_position).unwrap());
+                if &absolute > text_pos {
+                    let col = text_pos - curr;
+                    return Some(FilePosition {
+                        file: path.clone(),
+                        line: curr_line + 1,
+                        column: <TextSize as Into<usize>>::into(col) + 1,
+                    });
+                }
+                curr_position += line.len() + 1;
+                None
+            }
+            _ => None,
+        };
+        if let Some(pos) = pos_found {
+            return pos;
+        }
+    }
+    None
+}
+
+// Removes potential parenthesis wrapping the lambda
+// (( x: ... )) -> x: ...
+// Returns the lambda node, if one exists
+// Aborts if the node is not a lambda
+fn unpack_lambda(node: &SyntaxNode) -> Option<SyntaxNode> {
+    for ev in node.preorder() {
+        let res = match ev {
+            WalkEvent::Enter(node) => {
+                match node.kind() {
+                    // The top level callpackage lambda
+                    rnix::SyntaxKind::NODE_PAREN => None,
+                    rnix::SyntaxKind::NODE_LAMBDA => Some(node),
+                    _ => {
+                        println!(
+                            "Unexpected node kind: {:?}. Expected Parenthesis '(x: ...)' or Lambda 'x: ... '",
+                            node.kind()
+                        );
+                        exit(1);
+                    }
+                }
+            }
+            _ => None,
+        };
+        if let Some(res) = res {
+            return Some(res);
+        }
+    }
+    None
+}
+
+fn get_call_package_lambda_body(root: &SyntaxNode) -> Option<SyntaxNode> {
+    for ev in root.preorder() {
+        let res = match ev {
+            WalkEvent::Enter(node) => {
+                match node.kind() {
+                    // The top level callpackage lambda
+                    rnix::SyntaxKind::NODE_LAMBDA => node.last_child(),
+                    _ => None,
+                }
+            }
+            _ => None,
+        };
+        if let Some(res) = res {
+            return Some(res);
+        }
+    }
+    None
+}
+
+fn contains_make_overridable(apply_lhs: &SyntaxNode) -> bool {
+    for ev in apply_lhs.preorder() {
+        let res = match ev {
+            WalkEvent::Enter(node) => {
+                match node.kind() {
+                    // Potential apply of makeOverridable ()
+                    rnix::SyntaxKind::NODE_IDENT => match node.text().to_string().as_str() {
+                        "makeOverridable" => Some(node),
+                        _ => None,
+                    },
+                    _ => None,
+                }
+            }
+            _ => None,
+        };
+        if let Some(_res) = res {
+            return true;
+        }
+    }
+    false
+}
+
+fn get_apply_make_overridable_body(body: &SyntaxNode) -> Option<SyntaxNode> {
+    for ev in body.preorder() {
+        let res = match ev {
+            WalkEvent::Enter(node) => {
+                match node.kind() {
+                    // Potential apply of makeOverridable ()
+                    rnix::SyntaxKind::NODE_APPLY => {
+                        // TODO: check if the first child is actually contains makeOverridable
+                        if contains_make_overridable(&node.first_child().unwrap()) {
+                            node.last_child()
+                        } else {
+                            None
+                        }
+                    }
+                    _ => None,
+                }
+            }
+            _ => None,
+        };
+        if let Some(res) = res {
+            return Some(res);
+        }
+    }
+    None
+}
+
 // Take a list of lookup operations
 // Since iterating over the AST can be expensive
 fn init_node_idx(
@@ -164,6 +474,22 @@ impl<'a> DocComment<'a> for DocIndex<'a> {
         }
         return None;
     }
+
+    fn get_node_at_position(&self, position: &'a FilePosition) -> &Option<SyntaxNode> {
+        if &position.file != self.file {
+            println!(
+                "Invalid usage of get_node_at_position: File {:?} does not match index file source {:?}",
+                self.file, position.file
+            );
+            exit(1);
+        }
+        let res = self
+            .pos_idx
+            .get(&(position.line, position.column))
+            .map(|idx| self.node_idx.get(&idx))
+            .flatten();
+        return res.unwrap_or(&None);
+    }
 }
 
 #[derive(Debug)]