add end-of-line signature comments

5 files changed, 124 insertions, 64 deletions

diff --git a/changelog.md b/changelog.md
index 847bc9e7..75c92fae 100644
--- a/changelog.md
+++ b/changelog.md
@@ -73,6 +73,8 @@ This version is not yet released. If you are reading this on the website, then t
 - Remove previously deprecated function strands
 ### Interpreter
 - Add the `uiua find` command, which finds Uiua code that matches the given unformatted text
+- Add checking for end-of-line signature comments
+  - These are documented in the [Documenting Code](https://uiua.org/tutorial/documentation) tutorial
 - Some optimizations
   - Pervasive function machinery has been totally rewritten
     - Observed performance improvements of up to 12x
diff --git a/site/src/other_tutorial.rs b/site/src/other_tutorial.rs
index 9491f33f..6ecf13b1 100644
--- a/site/src/other_tutorial.rs
+++ b/site/src/other_tutorial.rs
@@ -211,6 +211,10 @@ pub fn Documentation() -> impl IntoView {
         <p>"The "<code>"?"</code>" is similar to the "<Prim prim=Primitive::Stack/>" function because the arguments indicate the intended state of the stack before the function is called."</p>
         <p>"If you also want to give names to a function's outputs, you can list them in front of the "<code>"?"</code>". This lets you read the comment signature right-to-left, the same way as normal Uiua code."</p>
         <Editor example="# Quotient Remainder ? Divisor Dividend\nDivRem ← ⌊⊃÷◿\nDivRem 3 7"/>
+        <p>"These kinds of comments can also be put at the end of lines. The signature of the line will be checked against the signature specified in the comment."</p>
+        <Editor example="1 2 # A B ?\n+⌵  # Sum ? A B\n⇡+5 # Res ? Foo Bar"/> // Should fail
+        <p>"These can be put in functions as well."</p>
+        <Editor example="# Get the average of a list\n# ? List\nAvg ← (\n  ⟜/+ # List Sum ? List\n  ⧻   # Length ? List\n  ÷   # Avg ? Length List\n)"/>
 
         <Hd id="track-caller"><code>"# Track caller!"</code></Hd>
 
diff --git a/src/assembly.rs b/src/assembly.rs
index 2a3be28f..94be0eee 100644
--- a/src/assembly.rs
+++ b/src/assembly.rs
@@ -1,4 +1,4 @@
-use std::{fmt, iter::once, path::PathBuf, sync::Arc};
+use std::{fmt, path::PathBuf, str::FromStr, sync::Arc};
 
 use dashmap::DashMap;
 use ecow::{eco_vec, EcoString, EcoVec};
@@ -518,6 +518,66 @@ pub struct DocCommentArg {
     pub ty: Option<EcoString>,
 }
 
+impl FromStr for DocCommentSig {
+    type Err = ();
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if s.trim_end().ends_with('?') && !s.trim_end().ends_with(" ?")
+            || !(s.chars()).all(|c| c.is_whitespace() || "?:".contains(c) || is_ident_char(c))
+        {
+            return Err(());
+        }
+        // Split into args and outputs
+        let (mut outputs_text, mut args_text) = s.split_once('?').ok_or(())?;
+        outputs_text = outputs_text.trim();
+        args_text = args_text.trim();
+        // Parse args and outputs
+        let mut args = Vec::new();
+        let mut outputs = Vec::new();
+        for (args, text) in [(&mut args, args_text), (&mut outputs, outputs_text)] {
+            // Tokenize text
+            let mut tokens = Vec::new();
+            for frag in text.split_whitespace() {
+                for (i, token) in frag.split(':').enumerate() {
+                    if i > 0 {
+                        tokens.push(":");
+                    }
+                    tokens.push(token);
+                }
+            }
+            // Parse tokens into args
+            let mut curr_arg_name = None;
+            let mut tokens = tokens.into_iter().peekable();
+            while let Some(token) = tokens.next() {
+                if token == ":" {
+                    let ty = tokens.next().unwrap_or_default();
+                    args.push(DocCommentArg {
+                        name: curr_arg_name.take().unwrap_or_default(),
+                        ty: if ty.is_empty() { None } else { Some(ty.into()) },
+                    });
+                } else {
+                    if let Some(curr) = curr_arg_name.take() {
+                        args.push(DocCommentArg {
+                            name: curr,
+                            ty: None,
+                        });
+                    }
+                    curr_arg_name = Some(token.into());
+                }
+            }
+            if let Some(curr) = curr_arg_name.take() {
+                args.push(DocCommentArg {
+                    name: curr,
+                    ty: None,
+                });
+            }
+        }
+        Ok(DocCommentSig {
+            args,
+            outputs: (!outputs.is_empty()).then_some(outputs),
+        })
+    }
+}
+
 impl From<String> for DocComment {
     fn from(text: String) -> Self {
         Self::from(text.as_str())
@@ -533,60 +593,10 @@ impl From<&str> for DocComment {
                 && (line.chars()).all(|c| c.is_whitespace() || "?:".contains(c) || is_ident_char(c))
         });
         let raw_text = if let Some(i) = sig_line {
-            let sig_text = text.lines().nth(i).unwrap();
-            // Split into args and outputs
-            let (mut outputs_text, mut args_text) = sig_text.split_once('?').unwrap();
-            outputs_text = outputs_text.trim();
-            args_text = args_text.trim();
-            // Parse args and outputs
-            let mut args = Vec::new();
-            let mut outputs = Vec::new();
-            for (args, text) in [(&mut args, args_text), (&mut outputs, outputs_text)] {
-                // Tokenize text
-                let mut tokens = Vec::new();
-                for frag in text.split_whitespace() {
-                    for (i, token) in frag.split(':').enumerate() {
-                        if i > 0 {
-                            tokens.push(":");
-                        }
-                        tokens.push(token);
-                    }
-                }
-                // Parse tokens into args
-                let mut curr_arg_name = None;
-                let mut tokens = tokens.into_iter().peekable();
-                while let Some(token) = tokens.next() {
-                    if token == ":" {
-                        let ty = tokens.next().unwrap_or_default();
-                        args.push(DocCommentArg {
-                            name: curr_arg_name.take().unwrap_or_default(),
-                            ty: if ty.is_empty() { None } else { Some(ty.into()) },
-                        });
-                    } else {
-                        if let Some(curr) = curr_arg_name.take() {
-                            args.push(DocCommentArg {
-                                name: curr,
-                                ty: None,
-                            });
-                        }
-                        curr_arg_name = Some(token.into());
-                    }
-                }
-                if let Some(curr) = curr_arg_name.take() {
-                    args.push(DocCommentArg {
-                        name: curr,
-                        ty: None,
-                    });
-                }
-            }
-
-            sig = Some(DocCommentSig {
-                args,
-                outputs: (!outputs.is_empty()).then_some(outputs),
-            });
+            sig = text.lines().nth(i).unwrap().parse().ok();
 
             let mut text: EcoString = (text.lines().take(i))
-                .chain(once("\n"))
+                .chain(["\n"])
                 .chain(text.lines().skip(i + 1))
                 .flat_map(|s| s.chars().chain(Some('\n')))
                 .collect();
diff --git a/src/compile/mod.rs b/src/compile/mod.rs
index 382003c8..09e3185b 100644
--- a/src/compile/mod.rs
+++ b/src/compile/mod.rs
@@ -35,10 +35,10 @@ use crate::{
     lsp::{CodeMeta, ImportSrc, SigDecl},
     optimize::{optimize_instrs, optimize_instrs_mut},
     parse::{count_placeholders, flip_unsplit_lines, parse, split_words},
-    Array, Assembly, BindingKind, Boxed, Diagnostic, DiagnosticKind, DocComment, GitTarget, Ident,
-    ImplPrimitive, InputSrc, IntoInputSrc, IntoSysBackend, Primitive, RunMode, SemanticComment,
-    SysBackend, Uiua, UiuaError, UiuaErrorKind, UiuaResult, Value, CONSTANTS, EXAMPLE_UA,
-    SUBSCRIPT_NUMS, VERSION,
+    Array, Assembly, BindingKind, Boxed, Diagnostic, DiagnosticKind, DocComment, DocCommentSig,
+    GitTarget, Ident, ImplPrimitive, InputSrc, IntoInputSrc, IntoSysBackend, Primitive, RunMode,
+    SemanticComment, SysBackend, Uiua, UiuaError, UiuaErrorKind, UiuaResult, Value, CONSTANTS,
+    EXAMPLE_UA, SUBSCRIPT_NUMS, VERSION,
 };
 
 /// The Uiua compiler
@@ -579,7 +579,7 @@ code:
         // Compile top-level words
 
         // Populate prelude
-        for line in &lines {
+        for line in &mut lines {
             let mut words = line.iter().filter(|w| !matches!(w.value, Word::Spaces));
             if words.clone().count() == 1 {
                 let word = words.next().unwrap();
@@ -591,6 +591,7 @@ code:
                         } else {
                             prelude.comment = Some(c.as_str().into());
                         }
+                        line.clear();
                     }
                     Word::SemanticComment(SemanticComment::NoInline) => {
                         prelude.flags |= FunctionFlags::NO_INLINE;
@@ -625,6 +626,7 @@ code:
                     Word::Char(_) | Word::Number(..) | Word::String(_) | Word::MultilineString(_)
                 )
             });
+            let line_sig_comment = line_sig(&line);
             // Compile the words
             let instr_count_before = self.asm.instrs.len();
             let binding_count_before = self.asm.bindings.len();
@@ -635,6 +637,16 @@ code:
             let mut line_eval_errored = false;
             match instrs_signature(&new_func.instrs) {
                 Ok(sig) => {
+                    // Check doc comment sig
+                    if let Some(comment_sig) = line_sig_comment {
+                        if !comment_sig.value.matches_sig(sig) {
+                            self.emit_diagnostic(
+                                format!("Line signature {sig} does not match comment"),
+                                DiagnosticKind::Warning,
+                                comment_sig.span.clone(),
+                            );
+                        }
+                    }
                     // Update scope stack height
                     if let Ok(height) = &mut self.scope.stack_height {
                         *height = (*height + sig.outputs).saturating_sub(sig.args);
@@ -909,6 +921,23 @@ code:
         }
         Ok(self.new_functions.pop().unwrap())
     }
+    // Compile a line, checking an end-of-line signature comment
+    fn compile_line(&mut self, line: Vec<Sp<Word>>, call: bool) -> UiuaResult<NewFunction> {
+        let comment_sig = line_sig(&line);
+        let new_func = self.compile_words(line, call)?;
+        if let Some(comment_sig) = comment_sig {
+            if let Ok(sig) = instrs_signature(&new_func.instrs) {
+                if !comment_sig.value.matches_sig(sig) {
+                    self.emit_diagnostic(
+                        format!("Line signature {sig} does not match comment"),
+                        DiagnosticKind::Warning,
+                        comment_sig.span.clone(),
+                    );
+                }
+            }
+        }
+        Ok(new_func)
+    }
     fn compile_operand_word(&mut self, word: Sp<Word>) -> UiuaResult<(NewFunction, Signature)> {
         let span = word.span.clone();
         let mut new_func = self.compile_words(vec![word], true)?;
@@ -1275,8 +1304,8 @@ code:
                 let mut inner = Vec::new();
                 let mut flags = FunctionFlags::default();
                 let line_count = arr.lines.len();
-                for lines in arr.lines.into_iter().rev() {
-                    let nf = self.compile_words(lines, true)?;
+                for line in arr.lines.into_iter().rev() {
+                    let nf = self.compile_line(line, true)?;
                     inner.extend(nf.instrs);
                     flags |= nf.flags;
                 }
@@ -1843,7 +1872,7 @@ code:
     ) -> UiuaResult<(FunctionId, Option<Signature>, NewFunction)> {
         let mut new_func = NewFunction::default();
         for line in func.lines {
-            let nf = self.compile_words(line, true)?;
+            let nf = self.compile_line(line, true)?;
             new_func.instrs.extend(nf.instrs);
             new_func.flags |= nf.flags;
         }
@@ -2730,3 +2759,20 @@ fn recurse_words_mut(words: &mut Vec<Sp<Word>>, f: &mut dyn FnMut(&mut Sp<Word>)
         }
     }
 }
+
+fn line_sig(line: &[Sp<Word>]) -> Option<Sp<DocCommentSig>> {
+    line.split_last()
+        .and_then(|(last, mut rest)| match &last.value {
+            Word::Comment(c) => c.parse::<DocCommentSig>().ok().map(|sig| {
+                while rest.last().is_some_and(|w| matches!(w.value, Word::Spaces)) {
+                    rest = &rest[..rest.len() - 1];
+                }
+                let span = (rest.first())
+                    .zip(rest.last())
+                    .map(|(f, l)| f.span.clone().merge(l.span.clone()))
+                    .unwrap_or_else(|| last.span.clone());
+                span.sp(sig)
+            }),
+            _ => None,
+        })
+}
diff --git a/todo.md b/todo.md
index 08f46a79..c35d24ec 100644
--- a/todo.md
+++ b/todo.md
@@ -1,9 +1,7 @@
 # Uiua Todo
 
 - 0.13
-  - End-of-line signature comments
-  - Align multi-line arrays in formatter
-  - Allow trailing newline in arrays
+  - Better testing
   - Make `un under` work...
     - `un` for `under` arg extensions
     - `un dip`