[CLI/PTB] Check for help flags after lexing (MystenLabs#16435)

## Description

Make help flag detection more robust by doing it after tokenisation, so
that the help flag is treated consistently like other `--commands` or
`-f`lags.

This PR also introduces a new notion of a `-f`lag to the lexer (which is
only used for `-h` at the moment).

## Test Plan

New unit tests:

```
crates/sui$ cargo nextest run -- lexer
```

Trying the command:

```
sui$ cargo run --bin sui -- --help
sui$ cargo run --bin sui -- -h
```

crates/sui/src/client_commands.rs

@@ -1,7 +1,7 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::client_ptb::ptb::{ptb_description, PTB};
+use crate::client_ptb::ptb::PTB;
 use std::{
     collections::{btree_map::Entry, BTreeMap},
     fmt::{Debug, Display, Formatter, Write},
@@ -1594,13 +1594,7 @@ impl SuiClientCommands {
                 SuiClientCommandResult::VerifySource
             }
             SuiClientCommands::PTB(ptb) => {
-                if ptb.args.contains(&"--help".to_string()) {
-                    ptb_description().print_long_help().unwrap();
-                } else if ptb.args.contains(&"-h".to_string()) || ptb.args.is_empty() {
-                    ptb_description().print_help().unwrap();
-                } else {
-                    ptb.execute(context).await?;
-                }
+                ptb.execute(context).await?;
                 SuiClientCommandResult::NoOutput
             }
         });

crates/sui/src/client_ptb/lexer.rs

@@ -279,12 +279,24 @@ impl<'l, I: Iterator<Item = &'l str>> Iterator for Lexer<'l, I> {
             }
 
             sp!(_, "-") => 'command: {
-                let Some(prefix) = self.eat_prefix("--") else {
+                self.bump();
+                let Some(next) = self.peek() else {
                     break 'command self.unexpected(c);
                 };
 
+                match next {
+                    sp!(_, "-") => {
+                        self.bump();
+                    }
+                    sp!(_, flag) if is_flag(flag.chars().next().unwrap()) => {
+                        self.bump();
+                        break 'command next.widen(c).map(|src| Lexeme(T::Flag, src));
+                    }
+                    sp!(_, _) => break 'command self.unexpected(next),
+                }
+
                 let Some(ident) = self.eat_while(is_ident_continue) else {
-                    break 'command self.unexpected(prefix);
+                    break 'command self.unexpected(c);
                 };
 
                 match ident {
@@ -297,7 +309,7 @@ impl<'l, I: Iterator<Item = &'l str>> Iterator for Lexer<'l, I> {
                             break 'command self.done(T::EarlyEof);
                         };
 
-                        file.widen(prefix).map(|src| Lexeme(T::Publish, src))
+                        file.widen(c).map(|src| Lexeme(T::Publish, src))
                     }
 
                     sp!(_, "upgrade") => {
@@ -309,10 +321,10 @@ impl<'l, I: Iterator<Item = &'l str>> Iterator for Lexer<'l, I> {
                             break 'command self.done(T::EarlyEof);
                         };
 
-                        file.widen(prefix).map(|src| Lexeme(T::Upgrade, src))
+                        file.widen(c).map(|src| Lexeme(T::Upgrade, src))
                     }
 
-                    sp!(_, _) => ident.widen(prefix).map(|src| Lexeme(T::Command, src)),
+                    sp!(_, _) => ident.widen(c).map(|src| Lexeme(T::Command, src)),
                 }
             }
 
@@ -321,6 +333,10 @@ impl<'l, I: Iterator<Item = &'l str>> Iterator for Lexer<'l, I> {
     }
 }
 
+fn is_flag(c: char) -> bool {
+    c.is_ascii_alphanumeric()
+}
+
 fn is_ident_start(c: char) -> bool {
     c.is_ascii_alphabetic() || c == '_'
 }
@@ -408,6 +424,20 @@ mod tests {
         insta::assert_debug_snapshot!(lex(addrs));
     }
 
+    #[test]
+    fn tokenize_commands() {
+        let cmds = vec!["--f00", "--Bar_baz", "--qux-quy"];
+
+        insta::assert_debug_snapshot!(lex(cmds));
+    }
+
+    #[test]
+    fn tokenize_flags() {
+        let flags = vec!["-h", "-a", "-Z", "-1"];
+
+        insta::assert_debug_snapshot!(lex(flags));
+    }
+
     #[test]
     fn tokenize_args() {
         let args = vec![

crates/sui/src/client_ptb/ptb.rs

@@ -1,11 +1,15 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::client_ptb::{
-    ast::{ParsedProgram, Program},
-    builder::PTBBuilder,
-    displays::Pretty,
-    error::{build_error_reports, PTBError},
+use crate::{
+    client_ptb::{
+        ast::{ParsedProgram, Program},
+        builder::PTBBuilder,
+        displays::Pretty,
+        error::{build_error_reports, PTBError},
+        token::{Lexeme, Token},
+    },
+    sp,
 };
 
 use anyhow::{anyhow, Error};
@@ -26,7 +30,7 @@ use sui_types::{
     transaction::{ProgrammableTransaction, Transaction, TransactionData},
 };
 
-use super::{ast::ProgramMetadata, parser::ProgramParser};
+use super::{ast::ProgramMetadata, lexer::Lexer, parser::ProgramParser};
 
 #[derive(Clone, Debug, Args)]
 #[clap(disable_help_flag = true)]
@@ -52,11 +56,21 @@ impl PTB {
     pub async fn execute(self, context: &mut WalletContext) -> Result<(), Error> {
         let source_string = to_source_string(self.args.clone());
 
-        let (program, program_metadata) = match crate::client_ptb::parser::ProgramParser::new(
-            self.args.iter().map(|s| s.as_str()),
-        )
-        .map_err(|e| vec![e])
-        .and_then(|parser| parser.parse())
+        // Tokenize once to detect help flags
+        let tokens = self.args.iter().map(|s| s.as_str());
+        for sp!(_, lexeme) in Lexer::new(tokens.clone()).into_iter().flatten() {
+            match lexeme {
+                Lexeme(Token::Command, "help") => return Ok(ptb_description().print_long_help()?),
+                Lexeme(Token::Flag, "h") => return Ok(ptb_description().print_help()?),
+                lexeme if lexeme.is_terminal() => break,
+                _ => continue,
+            }
+        }
+
+        // Tokenize and parse to get the program
+        let (program, program_metadata) = match ProgramParser::new(tokens)
+            .map_err(|e| vec![e])
+            .and_then(|parser| parser.parse())
         {
             Err(errors) => {
                 let suffix = if errors.len() > 1 { "s" } else { "" };
@@ -349,7 +363,7 @@ pub fn ptb_description() -> clap::Command {
             "Publish the Move package. It takes as input the folder where the package exists."
         ).long_help(
             "Publish the Move package. It takes as input the folder where the package exists.\
-            \n\nExample:\
+            \n\nExamples:\
             \n --move-call sui::tx_context::sender\
             \n --assign sender\
             \n --publish \".\"\
@@ -361,20 +375,20 @@ pub fn ptb_description() -> clap::Command {
             "Upgrade the move package. It takes as input the folder where the package exists."
         ).value_hint(ValueHint::DirPath))
         .arg(arg!(
-            --"preview" 
+            --"preview"
             "Preview the list of PTB transactions instead of executing them."
         ))
         .arg(arg!(
-            --"summary" 
+            --"summary"
             "Show only a short summary (digest, execution status, gas cost). \
             Do not use this flag when you need all the transaction data and the execution effects."
         ))
         .arg(arg!(
-            --"warn-shadows" 
+            --"warn-shadows"
             "Enable shadow warning when the same variable name is declared multiple times. Off by default."
         ))
         .arg(arg!(
-            --"json" 
+            --"json"
             "Return command outputs in json format."
         ))
 }

crates/sui/src/client_ptb/snapshots/sui__client_ptb__lexer__tests__tokenize_commands.snap

@@ -0,0 +1,46 @@
+---
+source: crates/sui/src/client_ptb/lexer.rs
+expression: lex(cmds)
+---
+[
+    Spanned {
+        span: Span {
+            start: 0,
+            end: 5,
+        },
+        value: Lexeme(
+            Command,
+            "f00",
+        ),
+    },
+    Spanned {
+        span: Span {
+            start: 6,
+            end: 15,
+        },
+        value: Lexeme(
+            Command,
+            "Bar_baz",
+        ),
+    },
+    Spanned {
+        span: Span {
+            start: 16,
+            end: 25,
+        },
+        value: Lexeme(
+            Command,
+            "qux-quy",
+        ),
+    },
+    Spanned {
+        span: Span {
+            start: 25,
+            end: 25,
+        },
+        value: Lexeme(
+            Eof,
+            "",
+        ),
+    },
+]

crates/sui/src/client_ptb/snapshots/sui__client_ptb__lexer__tests__tokenize_flags.snap

@@ -0,0 +1,56 @@
+---
+source: crates/sui/src/client_ptb/lexer.rs
+expression: lex(flags)
+---
+[
+    Spanned {
+        span: Span {
+            start: 0,
+            end: 2,
+        },
+        value: Lexeme(
+            Flag,
+            "h",
+        ),
+    },
+    Spanned {
+        span: Span {
+            start: 3,
+            end: 5,
+        },
+        value: Lexeme(
+            Flag,
+            "a",
+        ),
+    },
+    Spanned {
+        span: Span {
+            start: 6,
+            end: 8,
+        },
+        value: Lexeme(
+            Flag,
+            "Z",
+        ),
+    },
+    Spanned {
+        span: Span {
+            start: 9,
+            end: 11,
+        },
+        value: Lexeme(
+            Flag,
+            "1",
+        ),
+    },
+    Spanned {
+        span: Span {
+            start: 11,
+            end: 11,
+        },
+        value: Lexeme(
+            Eof,
+            "",
+        ),
+    },
+]

crates/sui/src/client_ptb/snapshots/sui__client_ptb__lexer__tests__unexpected_dash_dash.snap

@@ -1,16 +1,16 @@
 ---
 source: crates/sui/src/client_ptb/lexer.rs
-expression: tokens
+expression: lex(unexpected)
 ---
 [
     Spanned {
         span: Span {
             start: 0,
-            end: 2,
+            end: 1,
         },
         value: Lexeme(
             Unexpected,
-            "--",
+            "-",
         ),
     },
 ]

crates/sui/src/client_ptb/token.rs

@@ -15,6 +15,8 @@ pub struct Lexeme<'t>(
 pub enum Token {
     /// --[a-zA-Z0-9_-]+
     Command,
+    /// -[a-zA-Z0-9]
+    Flag,
     /// [a-zA-Z_][a-zA-Z0-9_-]*
     Ident,
     /// [1-9][0-9_]*
@@ -84,6 +86,7 @@ impl<'a> fmt::Display for Lexeme<'a> {
 
         match self.0 {
             T::Command => write!(f, "command '--{}'", self.1),
+            T::Flag => write!(f, "flag '-{}'", self.1),
             T::Ident => write!(f, "identifier '{}'", self.1),
             T::Number => write!(f, "number '{}'", self.1),
             T::HexNumber => write!(f, "hexadecimal number '0x{}'", self.1),
@@ -112,6 +115,7 @@ impl fmt::Display for Token {
         use Token as T;
         match self {
             T::Command => write!(f, "a command"),
+            T::Flag => write!(f, "a flag"),
             T::Ident => write!(f, "an identifier"),
             T::Number => write!(f, "a number"),
             T::HexNumber => write!(f, "a hexadecimal number"),