sway-fmt-v2 adds program type to the output

sway-fmt-v2/src/fmt.rs

@@ -1,5 +1,6 @@
 use crate::utils::{
     attributes::format_attributes, indent_style::Shape, newline_style::apply_newline_style,
+    program_type::insert_program_type,
 };
 use std::{path::Path, sync::Arc};
 use sway_core::BuildConfig;
@@ -38,8 +39,20 @@ impl Formatter {
         build_config: Option<&BuildConfig>,
     ) -> Result<FormattedCode, FormatterError> {
         let path = build_config.map(|build_config| build_config.canonical_root_module());
-        let items = sway_parse::parse_file(src, path)?.items;
-        let formatted_raw_newline = items
+        let module = sway_parse::parse_file(src, path)?;
+        // Get parsed items
+        let items = module.items;
+        // Get the program type (script, predicate, contract or library)
+        let program_type = module.kind;
+
+        // Formatted code will be pushed here with raw newline stlye.
+        // Which means newlines are not converted into system-specific versions by apply_newline_style
+        let mut raw_formatted_code = String::new();
+
+        // Insert program type to the formatted code.
+        insert_program_type(&mut raw_formatted_code, program_type);
+        // Insert parsed & formatted items into the formatted code.
+        raw_formatted_code += &items
             .into_iter()
             .map(|item| -> Result<String, FormatterError> {
                 use ItemKind::*;
@@ -60,12 +73,12 @@ impl Formatter {
             })
             .collect::<Result<Vec<String>, _>>()?
             .join("\n");
-        let mut formatted_code = String::from(&formatted_raw_newline);
+        let mut formatted_code = String::from(&raw_formatted_code);
         apply_newline_style(
             // The user's setting for `NewlineStyle`
             self.config.whitespace.newline_style,
             &mut formatted_code,
-            &formatted_raw_newline,
+            &raw_formatted_code,
         );
         Ok(formatted_code)
     }
@@ -93,8 +106,9 @@ enum Color {
                     Grey: (), }
         "#;
 
-        // Until #1995 is addressed we will not have contract; in the output
-        let correct_sway_code = r#"enum Color {
+        let correct_sway_code = r#"contract;
+
+enum Color {
  Blue : (),
  Green : (),
  Red : (),
@@ -117,8 +131,9 @@ enum Color {
                     Grey: (), }
         "#;
 
-        // Until #1995 is addressed we will not have contract; in the output
-        let correct_sway_code = r#"enum Color {
+        let correct_sway_code = r#"contract;
+
+enum Color {
  Blue   : (),
  Green  : (),
  Red    : (),

sway-fmt-v2/src/utils.rs

@@ -2,3 +2,4 @@ pub mod attributes;
 pub mod bracket;
 pub mod indent_style;
 pub mod newline_style;
+pub mod program_type;

sway-fmt-v2/src/utils/program_type.rs

@@ -0,0 +1,30 @@
+use sway_parse::ModuleKind;
+use sway_types::Spanned;
+
+/// Insert the program type without applying a formatting to it.
+///
+/// Possible list of program types:
+///     - Script
+///     - Contract
+///     - Predicate
+///     - Library
+pub(crate) fn insert_program_type(push_to: &mut String, module_kind: ModuleKind) {
+    match module_kind {
+        ModuleKind::Script { script_token } => push_to.push_str(script_token.span().as_str()),
+        ModuleKind::Contract { contract_token } => push_to.push_str(contract_token.span().as_str()),
+        ModuleKind::Predicate { predicate_token } => {
+            push_to.push_str(predicate_token.span().as_str())
+        }
+        ModuleKind::Library {
+            library_token,
+            name,
+        } => {
+            push_to.push_str(library_token.span().as_str());
+            push_to.push(' ');
+            push_to.push_str(name.as_str());
+        }
+    };
+    push_to.push(';');
+    push_to.push('\n');
+    push_to.push('\n');
+}