diff --git a/src/helpers/testCode.ts b/src/helpers/testCode.ts
index 89d47ee..398d281 100644
--- a/src/helpers/testCode.ts
+++ b/src/helpers/testCode.ts
@@ -105,26 +105,69 @@ function shell(command, options) {
     });
 }
 
+function removeMultilineComments(data) {
+    return data.replace(/{-[\s\S]*?-}/g, (comment) => {
+        return comment.replace(/\S/g, ' ');
+    });
+}
+
 function removeMainFunction(data) {
-    const dataArray = data.toString().split('\n');
+    const decl_regex = /^main\s*::.*$/m;
+    const def_regex = /^main\s*=.*$/m;
+    const comment_regex = /^\s*--.*$/m;
+    const other_regex = /^\S+.*$/m;
+    const dataArray = removeMultilineComments(data).split('\n');
+
+    let decl_start;
+    let decl_end;
+    let def_start;
+    let def_end;
 
-    let start;
-    let end;
     for (let i = 0; i < dataArray.length; i++) {
-        if (dataArray[i].slice(0, 6) === 'main =') {
-            start = i;
-            end = i;
-        } else if (start !== undefined) {
-            if (dataArray[i] === '') {
-                break;
-            } else {
-                end++;
-            }   
+    	//main type signature
+    	if (decl_regex.test(dataArray[i])) {
+            if (def_start !== undefined && def_end === undefined) {
+                def_end = i;
+            }
+            if (decl_start === undefined) {
+                decl_start = i;
+            }
+        }
+        //main definition
+        else if (def_regex.test(dataArray[i])) {
+            if (decl_start !== undefined && decl_end === undefined) {
+                decl_end = i;
+            }
+            if (def_start === undefined) {
+                def_start = i;
+            }
         }
+        //skip comments
+        else if (comment_regex.test(dataArray[i])) {
+        	//do nothing
+        }
+        //different expression/declaration
+        else if (other_regex.test(dataArray[i])) {
+            if (decl_start !== undefined && decl_end === undefined) {
+                decl_end = i;
+            }
+            if (def_start !== undefined && def_end === undefined) {
+                def_end = i;
+            }
+        }
+    }
+
+    if (decl_start !== undefined && decl_end === undefined) {
+        decl_end = dataArray.length;
+    }
+    if (def_start !== undefined && def_end === undefined) {
+        def_end = dataArray.length;
     }
 
-    dataArray.splice(start, end - start + 1);
-    return dataArray.join('\n');
+    return dataArray.filter((value, index, arr) => {
+        return (decl_start === undefined || index < decl_start || index >= decl_end) &&
+               (def_start === undefined || index < def_start || index >= def_end);
+    }).join('\n');
 }
 
 export function testHaskellFile(filePath, stackWd) {
@@ -133,24 +176,27 @@ export function testHaskellFile(filePath, stackWd) {
 
         // Not Stack project
         if (stackWd === undefined) {
-            fs.createReadStream(filePath).pipe(fs.createWriteStream(newPath));
-            fs.readFile(newPath, 'utf-8', (err, data) => {
-                if (err) reject(err);
-
-                const newValue = '{-# LANGUAGE TemplateHaskell #-}\nimport Test.QuickCheck.All\n' + removeMainFunction(data)
-                    + '\nreturn []\nrunTests = $quickCheckAll\nmain = runTests';
-
-                fs.writeFile(newPath, newValue, 'utf-8', err => {
+            let writer = fs.createWriteStream(newPath);
+            fs.createReadStream(filePath).pipe(writer);
+            writer.once('finish', () => {
+                fs.readFile(newPath, 'utf-8', (err, data) => {
                     if (err) reject(err);
-                    console.log('QuickCheking...');
-                    
-                    shell(`stack runhaskell "${newPath}"`, {}).then(std => {
-                        console.log(std[0]);
-                        fs.unlinkSync(newPath);
-                        resolve(parseStdout(std[0]));
-                    }).catch(error => {
-                        fs.unlinkSync(newPath);
-                        reject(error);
+
+                    const newValue = '{-# LANGUAGE TemplateHaskell #-}\nimport Test.QuickCheck.All\n' + removeMainFunction(data)
+                        + '\nreturn []\nrunTests = $quickCheckAll\nmain = runTests';
+
+                    fs.writeFile(newPath, newValue, 'utf-8', err => {
+                        if (err) reject(err);
+                        console.log('QuickCheking...');
+                        
+                        shell(`stack runhaskell "${newPath}"`, {}).then(std => {
+                            console.log(std[0]);
+                            fs.unlinkSync(newPath);
+                            resolve(parseStdout(std[0]));
+                        }).catch(error => {
+                            fs.unlinkSync(newPath);
+                            reject(error);
+                        });
                     });
                 });
             });