sway-fmt-v2 adds program type to the output

sway-fmt-v2/src/fmt.rs

@@ -1,4 +1,9 @@
-use crate::utils::indent_style::Shape;
+use crate::{
+    config::whitespace::NewlineStyle,
+    utils::{
+        indent_style::Shape, newline_style::apply_newline_style, program_type::insert_program_type,
+    },
+};
 use std::{path::Path, sync::Arc};
 use sway_core::BuildConfig;
 use sway_parse::ItemKind;
@@ -37,8 +42,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 formatted_raw_newline = String::new();
+
+        // Insert program type to the formatted code.
+        insert_program_type(&mut formatted_raw_newline, program_type);
+        // Insert parsed & formatted items into the formatted code.
+        formatted_raw_newline += &items
             .into_iter()
             .map(|item| -> Result<String, FormatterError> {
                 use ItemKind::*;

sway-fmt-v2/src/utils.rs

@@ -1,3 +1,4 @@
 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,29 @@
+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');
+}