diff --git a/src/header.rs b/src/header.rs
index 3b2c5f4..eec2980 100644
--- a/src/header.rs
+++ b/src/header.rs
@@ -11,6 +11,20 @@ pub struct HeaderLine<'ln> {
     pub dejagnu_header: String,
 }
 
+pub fn parse_additional_options(code: &str) -> Vec<HeaderLine> {
+    let mut headers = Vec::new();
+
+    for (line_number, line) in code.lines().enumerate() {
+        let line = line.trim();
+        if line.is_empty() || line.starts_with("fn") || line.starts_with("mod") {
+            continue;
+        }
+        if line.trim_start().starts_with("//@") {
+            headers.push(add_additional_options(line, line_number).unwrap());
+        }
+    }
+    headers
+}
 
 fn add_additional_options(code: &str, line_number: usize) -> Option<HeaderLine> {
     //TODO: If we know the file extension, then update this to
diff --git a/src/transform.rs b/src/transform.rs
index ab9c213..308727a 100644
--- a/src/transform.rs
+++ b/src/transform.rs
@@ -1,7 +1,7 @@
 //! This module contains the code transformation logic.
 
 use {
-    crate::{errors, regex},
+    crate::{errors, header::parse_additional_options, regex},
     anyhow::Result,
 };
 
@@ -11,40 +11,52 @@ pub fn transform_code(code: &str, stderr_file: Option<&str>) -> Result<String> {
     let errors = errors::load_error(code, stderr_file);
     // For storing the transformed code
     let mut new_code = String::new();
+    let additional_options = parse_additional_options(code);
 
     let mut line_num = 1;
     // finding the respective line number and adding the error code
     for line in code.lines() {
         let mut new_line = line.to_string();
-        // TODO: This is not the efficient way to find respective line number
-        for error in errors.iter() {
-            // Checking the original line number
-            if (error.line_num as i32 - error.relative_line_num) != line_num {
-                continue;
+
+        if line.trim_start().starts_with("//@") {
+            for header in additional_options.iter() {
+                if header.line_number != line_num {
+                    continue;
+                }
+                new_line = header.dejagnu_header.to_string();
+                break;
             }
-            // In rustc test suites, the error directive is
-            // on the same line or on the next line, but not on the previous line
-            // See this: https://rustc-dev-guide.rust-lang.org/tests/ui.html#error-annotations
-            // For the error on the next line
-            if error.relative_line_num != 0 {
-                // We simply add the error message, not to worry about the code
-                // The error was printed by our overloaded `Display` trait
-                new_line = format!("{}", error);
-            } else {
-                // For the error on the same line, we need to add error message at the end of the line
-                let captures = regex!(r"//(?:\[(?P<revs>[\w\-,]+)])?~(?P<adjust>\||\^*)")
-                    .captures(line)
-                    .expect("Could not find the error directive");
+        } else {
+            // TODO: This is not the efficient way to find respective line number
+            for error in errors.iter() {
+                // Checking the original line number
+                if (error.line_num as i32 - error.relative_line_num) != line_num as i32 {
+                    continue;
+                }
+                // In rustc test suites, the error directive is
+                // on the same line or on the next line, but not on the previous line
+                // See this: https://rustc-dev-guide.rust-lang.org/tests/ui.html#error-annotations
+                // For the error on the next line
+                if error.relative_line_num != 0 {
+                    // We simply add the error message, not to worry about the code
+                    // The error was printed by our overloaded `Display` trait
+                    new_line = format!("{}", error);
+                } else {
+                    // For the error on the same line, we need to add error message at the end of the line
+                    let captures = regex!(r"//(?:\[(?P<revs>[\w\-,]+)])?~(?P<adjust>\||\^*)")
+                        .captures(line)
+                        .expect("Could not find the error directive");
 
-                // Get the part of comment before the sigil (e.g. `~^` or ~|)
-                let whole_match = captures.get(0).unwrap();
-                // Get the existing source code before the error directive //~ ERROR or similar to this
-                let before_match = &line[..whole_match.start()];
+                    // Get the part of comment before the sigil (e.g. `~^` or ~|)
+                    let whole_match = captures.get(0).unwrap();
+                    // Get the existing source code before the error directive //~ ERROR or similar to this
+                    let before_match = &line[..whole_match.start()];
 
-                // The error was printed by our overloaded `Display` trait
-                new_line = format!("{}{}", before_match, error);
+                    // The error was printed by our overloaded `Display` trait
+                    new_line = format!("{}{}", before_match, error);
+                }
+                break;
             }
-            break;
         }
         new_code.push_str(&new_line);
         new_code.push('\n');