Return the lexing stage produced by sway_parse from `forc-pkg::chec…

…k()` (FuelLabs#3675)

Previous to this PR, `forc_pkg::check()` would only return the
`ParseProgram` and `TyProgram` AST's. In the language server, we now
also need access to the initial lexing stage produced by `sway_parse` in
order to get access to the sway keyword tokens.

This PR now returns the lexing stage of a program along with the parsed
and typed stages.

closes FuelLabs#3524

forc-pkg/src/pkg.rs

@@ -31,16 +31,14 @@ use sway_core::{
         fuel_crypto,
         fuel_tx::{self, Contract, ContractId, StorageSlot},
     },
-    Engines, TypeEngine,
-};
-use sway_core::{
     language::{
+        lexed::LexedProgram,
         parsed::{ParseProgram, TreeType},
         ty,
     },
     semantic_analysis::namespace,
     source_map::SourceMap,
-    CompileResult, CompiledBytecode, FinalizedEntry,
+    CompileResult, CompiledBytecode, Engines, FinalizedEntry, TypeEngine,
 };
 use sway_error::error::CompileError;
 use sway_types::Ident;
@@ -2799,17 +2797,35 @@ fn update_json_type_declaration(
     }
 }
 
-/// A `CompileResult` thats type is a tuple containing a `ParseProgram` and `Option<ty::TyProgram>`
-type ParseAndTypedPrograms = CompileResult<(ParseProgram, Option<ty::TyProgram>)>;
+/// Contains the lexed, parsed, and typed compilation stages of a program.
+pub struct Programs {
+    pub lexed: LexedProgram,
+    pub parsed: ParseProgram,
+    pub typed: Option<ty::TyProgram>,
+}
+
+impl Programs {
+    pub fn new(
+        lexed: LexedProgram,
+        parsed: ParseProgram,
+        typed: Option<ty::TyProgram>,
+    ) -> Programs {
+        Programs {
+            lexed,
+            parsed,
+            typed,
+        }
+    }
+}
 
-/// Compile the entire forc package and return the parse and typed programs
+/// Compile the entire forc package and return the lexed, parsed and typed programs
 /// of the dependancies and project.
 /// The final item in the returned vector is the project.
 pub fn check(
     plan: &BuildPlan,
     terse_mode: bool,
     engines: Engines<'_>,
-) -> anyhow::Result<Vec<ParseAndTypedPrograms>> {
+) -> anyhow::Result<Vec<CompileResult<Programs>>> {
     let mut lib_namespace_map = Default::default();
     let mut source_map = SourceMap::new();
     // During `check`, we don't compile so this stays empty.
@@ -2829,27 +2845,28 @@ pub fn check(
             engines,
         )
         .expect("failed to create dependency namespace");
+
         let CompileResult {
             value,
             mut warnings,
             mut errors,
         } = parse(manifest, terse_mode, engines)?;
 
-        let parse_program = match value {
+        let (lexed, parsed) = match value {
             None => {
                 results.push(CompileResult::new(None, warnings, errors));
                 return Ok(results);
             }
-            Some(program) => program,
+            Some(modules) => modules,
         };
 
-        let ast_result = sway_core::parsed_to_ast(engines, &parse_program, dep_namespace, None);
+        let ast_result = sway_core::parsed_to_ast(engines, &parsed, dep_namespace, None);
         warnings.extend(ast_result.warnings);
         errors.extend(ast_result.errors);
 
         let typed_program = match ast_result.value {
             None => {
-                let value = Some((parse_program, None));
+                let value = Some(Programs::new(lexed, parsed, None));
                 results.push(CompileResult::new(value, warnings, errors));
                 return Ok(results);
             }
@@ -2862,7 +2879,7 @@ pub fn check(
 
         source_map.insert_dependency(manifest.dir());
 
-        let value = Some((parse_program, Some(typed_program)));
+        let value = Some(Programs::new(lexed, parsed, Some(typed_program)));
         results.push(CompileResult::new(value, warnings, errors));
     }
 
@@ -2878,7 +2895,7 @@ pub fn parse(
     manifest: &PackageManifestFile,
     terse_mode: bool,
     engines: Engines<'_>,
-) -> anyhow::Result<CompileResult<ParseProgram>> {
+) -> anyhow::Result<CompileResult<(LexedProgram, ParseProgram)>> {
     let profile = BuildProfile {
         terse: terse_mode,
         ..BuildProfile::debug()

forc-plugins/forc-doc/src/main.rs

@@ -64,8 +64,9 @@ pub fn main() -> Result<()> {
     let typed_program = match pkg::check(&plan, silent, engines)?
         .pop()
         .and_then(|compilation| compilation.value)
+        .and_then(|programs| programs.typed)
     {
-        Some((_, Some(typed_program))) => typed_program,
+        Some(typed_program) => typed_program,
         _ => bail!("CompileResult returned None"),
     };
     let raw_docs: Documentation = Document::from_ty_program(

forc/src/ops/forc_check.rs

@@ -28,6 +28,6 @@ pub fn check(command: CheckCommand, engines: Engines<'_>) -> Result<CompileResul
     let res = v
         .pop()
         .expect("there is guaranteed to be at least one elem in the vector")
-        .flat_map(|(_, tp)| CompileResult::new(tp, vec![], vec![]));
+        .flat_map(|programs| CompileResult::new(programs.typed, vec![], vec![]));
     Ok(res)
 }

sway-ast/src/module.rs

@@ -1,5 +1,6 @@
 use crate::priv_prelude::*;
 
+#[derive(Clone, Debug)]
 pub struct Module {
     pub kind: ModuleKind,
     pub semicolon_token: SemicolonToken,
@@ -30,6 +31,7 @@ impl Spanned for Module {
     }
 }
 
+#[derive(Clone, Debug)]
 pub enum ModuleKind {
     Script {
         script_token: ScriptToken,

sway-core/src/language/lexed/mod.rs

@@ -0,0 +1,25 @@
+mod program;
+
+use crate::language::DepName;
+pub use program::LexedProgram;
+use sway_ast::Module;
+use sway_types::Ident;
+
+/// A module and its submodules in the form of a tree.
+#[derive(Debug, Clone)]
+pub struct LexedModule {
+    /// The content of this module in the form of a [Module].
+    pub module: Module,
+    /// Submodules introduced within this module using the `dep` syntax in order of declaration.
+    pub submodules: Vec<(DepName, LexedSubmodule)>,
+}
+
+/// A library module that was declared as a `dep` of another module.
+///
+/// Only submodules are guaranteed to be a `library` and have a `library_name`.
+#[derive(Debug, Clone)]
+pub struct LexedSubmodule {
+    /// The name of a submodule, parsed from the `library` declaration within the module itself.
+    pub library_name: Ident,
+    pub module: LexedModule,
+}

sway-core/src/language/lexed/program.rs

@@ -0,0 +1,17 @@
+use super::LexedModule;
+use crate::language::parsed::TreeType;
+
+/// A lexed, but not yet parsed or type-checked, Sway program.
+///
+/// Includes all modules in the form of a [LexedModule] tree accessed via the `root`.
+#[derive(Debug, Clone)]
+pub struct LexedProgram {
+    pub kind: TreeType,
+    pub root: LexedModule,
+}
+
+impl LexedProgram {
+    pub fn new(kind: TreeType, root: LexedModule) -> LexedProgram {
+        LexedProgram { kind, root }
+    }
+}

sway-core/src/language/mod.rs

@@ -2,6 +2,7 @@ mod asm;
 mod call_path;
 mod inline;
 mod lazy_op;
+pub mod lexed;
 mod literal;
 mod module;
 pub mod parsed;

sway-core/src/lib.rs

@@ -40,7 +40,7 @@ use sway_error::warning::CompileWarning;
 use sway_types::{ident::Ident, span, Spanned};
 pub use type_system::*;
 
-use language::{parsed, ty};
+use language::{lexed, parsed, ty};
 use transform::to_parsed_lang::{self, convert_module_kind};
 
 pub mod fuel_prelude {
@@ -49,7 +49,7 @@ pub mod fuel_prelude {
 
 pub use engine_threading::Engines;
 
-/// Given an input `Arc<str>` and an optional [BuildConfig], parse the input into a [SwayParseTree].
+/// Given an input `Arc<str>` and an optional [BuildConfig], parse the input into a [lexed::LexedProgram] and [parsed::ParseProgram].
 ///
 /// # Example
 /// ```ignore
@@ -66,13 +66,21 @@ pub fn parse(
     input: Arc<str>,
     engines: Engines<'_>,
     config: Option<&BuildConfig>,
-) -> CompileResult<parsed::ParseProgram> {
+) -> CompileResult<(lexed::LexedProgram, parsed::ParseProgram)> {
     CompileResult::with_handler(|h| match config {
         None => parse_in_memory(h, engines, input),
         // When a `BuildConfig` is given,
         // the module source may declare `dep`s that must be parsed from other files.
-        Some(config) => parse_module_tree(h, engines, input, config.canonical_root_module())
-            .map(|(kind, root)| parsed::ParseProgram { kind, root }),
+        Some(config) => parse_module_tree(h, engines, input, config.canonical_root_module()).map(
+            |(kind, lexed, parsed)| {
+                let lexed = lexed::LexedProgram {
+                    kind: kind.clone(),
+                    root: lexed,
+                };
+                let parsed = parsed::ParseProgram { kind, root: parsed };
+                (lexed, parsed)
+            },
+        ),
     })
 }
 
@@ -90,12 +98,25 @@ fn parse_in_memory(
     handler: &Handler,
     engines: Engines<'_>,
     src: Arc<str>,
-) -> Result<parsed::ParseProgram, ErrorEmitted> {
+) -> Result<(lexed::LexedProgram, parsed::ParseProgram), ErrorEmitted> {
     let module = sway_parse::parse_file(handler, src, None)?;
-    let (kind, tree) = to_parsed_lang::convert_parse_tree(handler, engines, module)?;
+    let (kind, tree) = to_parsed_lang::convert_parse_tree(handler, engines, module.clone())?;
     let submodules = Default::default();
     let root = parsed::ParseModule { tree, submodules };
-    Ok(parsed::ParseProgram { kind, root })
+    let lexed_program = lexed::LexedProgram::new(
+        kind.clone(),
+        lexed::LexedModule {
+            module,
+            submodules: Default::default(),
+        },
+    );
+    Ok((lexed_program, parsed::ParseProgram { kind, root }))
+}
+
+/// Contains the lexed and parsed submodules 'deps' of a module.
+struct Submodules {
+    lexed: Vec<(Ident, lexed::LexedSubmodule)>,
+    parsed: Vec<(Ident, parsed::ParseSubmodule)>,
 }
 
 /// Parse all dependencies `deps` as submodules.
@@ -104,9 +125,10 @@ fn parse_submodules(
     engines: Engines<'_>,
     module: &sway_ast::Module,
     module_dir: &Path,
-) -> Vec<(Ident, parsed::ParseSubmodule)> {
+) -> Submodules {
     // Assume the happy path, so there'll be as many submodules as dependencies, but no more.
-    let mut submods = Vec::with_capacity(module.dependencies().count());
+    let mut lexed_submods = Vec::with_capacity(module.dependencies().count());
+    let mut parsed_submods = Vec::with_capacity(lexed_submods.capacity());
 
     module.dependencies().for_each(|dep| {
         // Read the source code from the dependency.
@@ -124,7 +146,7 @@ fn parse_submodules(
             }
         };
 
-        if let Ok((kind, module)) =
+        if let Ok((kind, lexed_module, parse_module)) =
             parse_module_tree(handler, engines, dep_str.clone(), dep_path.clone())
         {
             let library_name = match kind {
@@ -140,15 +162,23 @@ fn parse_submodules(
             // is where we should use it.
             let dep_alias = None;
             let dep_name = dep_alias.unwrap_or_else(|| library_name.clone());
-            let submodule = parsed::ParseSubmodule {
+            let parse_submodule = parsed::ParseSubmodule {
+                library_name: library_name.clone(),
+                module: parse_module,
+            };
+            let lexed_submodule = lexed::LexedSubmodule {
                 library_name,
-                module,
+                module: lexed_module,
             };
-            submods.push((dep_name, submodule));
+            lexed_submods.push((dep_name.clone(), lexed_submodule));
+            parsed_submods.push((dep_name, parse_submodule));
         }
     });
 
-    submods
+    Submodules {
+        lexed: lexed_submods,
+        parsed: parsed_submods,
+    }
 }
 
 /// Given the source of the module along with its path,
@@ -158,7 +188,7 @@ fn parse_module_tree(
     engines: Engines<'_>,
     src: Arc<str>,
     path: Arc<PathBuf>,
-) -> Result<(parsed::TreeType, parsed::ParseModule), ErrorEmitted> {
+) -> Result<(parsed::TreeType, lexed::LexedModule, parsed::ParseModule), ErrorEmitted> {
     // Parse this module first.
     let module_dir = path.parent().expect("module file has no parent directory");
     let module = sway_parse::parse_file(handler, src, Some(path.clone()))?;
@@ -168,9 +198,17 @@ fn parse_module_tree(
     let submodules = parse_submodules(handler, engines, &module, module_dir);
 
     // Convert from the raw parsed module to the `ParseTree` ready for type-check.
-    let (kind, tree) = to_parsed_lang::convert_parse_tree(handler, engines, module)?;
+    let (kind, tree) = to_parsed_lang::convert_parse_tree(handler, engines, module.clone())?;
 
-    Ok((kind, parsed::ParseModule { tree, submodules }))
+    let lexed = lexed::LexedModule {
+        module,
+        submodules: submodules.lexed,
+    };
+    let parsed = parsed::ParseModule {
+        tree,
+        submodules: submodules.parsed,
+    };
+    Ok((kind, lexed, parsed))
 }
 
 fn module_path(parent_module_dir: &Path, dep: &sway_ast::Dependency) -> PathBuf {
@@ -308,7 +346,7 @@ pub fn compile_to_ast(
         mut warnings,
         mut errors,
     } = parse(input, engines, build_config);
-    let parse_program = match parse_program_opt {
+    let (.., parse_program) = match parse_program_opt {
         Some(parse_program) => parse_program,
         None => return deduped_err(warnings, errors),
     };
@@ -910,7 +948,7 @@ fn test_unary_ordering() {
     );
     let mut warnings: Vec<CompileWarning> = Vec::new();
     let mut errors: Vec<CompileError> = Vec::new();
-    let prog = prog.unwrap(&mut warnings, &mut errors);
+    let (.., prog) = prog.unwrap(&mut warnings, &mut errors);
     // this should parse as `(!a) && b`, not `!(a && b)`. So, the top level
     // expression should be `&&`
     if let parsed::AstNode {

sway-lsp/src/core/session.rs

@@ -14,7 +14,7 @@ use crate::{
 use dashmap::DashMap;
 use forc_pkg::{self as pkg};
 use parking_lot::RwLock;
-use pkg::manifest::ManifestFile;
+use pkg::{manifest::ManifestFile, Programs};
 use std::{fs::File, io::Write, path::PathBuf, sync::Arc};
 use sway_core::{
     declaration_engine::DeclarationEngine,
@@ -151,14 +151,16 @@ impl Session {
                 errors,
             } = res;
 
-            // FIXME(Centril): Refactor parse_ast_to_tokens + parse_ast_to_typed_tokens
-            // due to the new API.g
-            let (parsed, typed) = match value {
-                None => (None, None),
-                Some((pp, tp)) => (Some(pp), tp),
-            };
+            if value.is_none() {
+                continue;
+            }
+            let Programs {
+                lexed: _, //(Josh): use this to collect keyword & bracket tokens
+                parsed,
+                typed,
+            } = value.unwrap();
 
-            let parsed_res = CompileResult::new(parsed, warnings.clone(), errors.clone());
+            let parsed_res = CompileResult::new(Some(parsed), warnings.clone(), errors.clone());
             let ast_res = CompileResult::new(typed, warnings, errors);
 
             let parse_program = self.compile_res_to_parse_program(&parsed_res)?;