1 files changed, 43 insertions, 20 deletions

diff --git a/lix-doc/src/pprint.rs b/lix-doc/src/pprint.rs
index 7e73d2d20..2a72c4069 100644
--- a/lix-doc/src/pprint.rs
+++ b/lix-doc/src/pprint.rs
@@ -1,36 +1,59 @@
 // SPDX-FileCopyrightText: 2024 Jade Lovelace
-//
+// SPDX-FileCopyrightText: 2024 Lunaphied
 // SPDX-License-Identifier: BSD-2-Clause OR MIT
 
-use rnix::types::{Lambda, TypedNode};
-use rnix::SyntaxKind::*;
+use rnix::ast::{Expr, Lambda};
+use rowan::ast::AstNode;
 
 /// Pretty-prints the arguments to a function
 pub fn pprint_args(lambda: &Lambda) -> String {
     // TODO: handle docs directly on NODE_IDENT args (uncommon case)
     let mut lambda = lambda.clone();
+    let mut depth = 0;
     let mut out = String::new();
     loop {
-        let arg = lambda.arg().unwrap();
-        match arg.kind() {
-            NODE_IDENT => {
-                out += &format!("*{}*", &arg.to_string());
-                out.push_str(": ");
-                let body = lambda.body().unwrap();
-                if body.kind() == NODE_LAMBDA {
-                    lambda = Lambda::cast(body).unwrap();
-                } else {
-                    break;
+        let arg = lambda.param().unwrap();
+        for child in arg.syntax().children_with_tokens() {
+            //dbg!(child.kind());
+            match child {
+                rowan::NodeOrToken::Node(node) => {
+                    out.push_str(&node.text().to_string());
+                    if node.kind() == rnix::SyntaxKind::NODE_PAT_ENTRY {
+                        out.push_str(&",\n");
+                    }
+                }
+                rowan::NodeOrToken::Token(token) => {
+                    use rnix::SyntaxKind::{
+                        TOKEN_COMMENT, TOKEN_ELLIPSIS, TOKEN_L_BRACE, TOKEN_QUESTION, TOKEN_R_BRACE,
+                    };
+                    match token.kind() {
+                        TOKEN_COMMENT | TOKEN_ELLIPSIS | TOKEN_QUESTION | TOKEN_L_BRACE
+                        | TOKEN_R_BRACE => {
+                            //dbg!(&token);
+                            out.push_str(&token.text().to_string());
+                            if token.kind() == TOKEN_COMMENT {
+                                out.push('\n');
+                            }
+                        }
+                        _ => {}
+                    }
+                    //out.push_str(&token.text().to_string());
                 }
             }
-            NODE_PATTERN => {
-                out += &format!("*{}*", &arg.to_string());
-                out.push_str(": ");
-                break;
-            }
-            t => {
-                unreachable!("unhandled arg type {:?}", t);
+        }
+        out.push_str(": ");
+        let body = lambda.body().unwrap();
+        if let Expr::Lambda(inner) = body {
+            lambda = inner;
+            // If we recurse we want the next line of recursion to be indented and on a new line.
+            out.push('\n');
+            for _ in 0..=depth {
+                out.push('\t');
             }
+            depth += 1;
+        } else {
+            // If we don't find an inner lambda we're done with argument handling.
+            break;
         }
     }
     out.push_str("...");