1 files changed, 24 insertions, 15 deletions

diff --git a/src/index.js b/src/index.js
index e09d247..72da4d0 100644
--- a/src/index.js
+++ b/src/index.js
@@ -10,7 +10,8 @@ const spawn = require('child_process').spawn
 const path = require('path')
 const retryPromise = require('promise-retry')
 
-const psModuleRegex = /(?:^|\n)module\s+([\w\.]+)/i
+const ffiModuleRegex = /\/\/\s+module\s+([\w\.]+)/i
+const srcModuleRegex = /(?:^|\n)module\s+([\w\.]+)/i
 const requireRegex = /require\(['"]\.\.\/([\w\.]+)['"]\)/g
 
 module.exports = function purescriptLoader(source, map) {
@@ -77,7 +78,7 @@ module.exports = function purescriptLoader(source, map) {
     })
   }
 
-  const psModuleName = match(psModuleRegex, source)
+  const psModuleName = match(srcModuleRegex, source)
   const psModule = {
     name: psModuleName,
     load: js => callback(null, js),
@@ -135,7 +136,7 @@ function toJavaScript(psModule) {
 
   return Promise.props({
     js: fs.readFileAsync(jsPath, 'utf8'),
-    psModuleMap: psModuleMap(options.src, cache)
+    psModuleMap: psModuleMap(options, cache)
   }).then(result => {
     let js = ''
 
@@ -147,15 +148,13 @@ function toJavaScript(psModule) {
     } else {
       // replace require paths to output files generated by psc with paths
       // to purescript sources, which are then also run through this loader.
-      const foreignRequire = 'require("' + path.resolve(
-        path.join(psModule.options.output, psModule.name, 'foreign.js')
-      ) + '")'
-
       js = result.js
         .replace(requireRegex, (m, p1) => {
-          return 'require("' + result.psModuleMap[p1] + '")'
+          return 'require("' + result.psModuleMap[p1].src + '")'
+        })
+        .replace(/require\(['"]\.\/foreign['"]\)/g, (m, p1) => {
+          return 'require("' + result.psModuleMap[psModule.name].ffi + '")'
         })
-        .replace(/require\(['"]\.\/foreign['"]\)/g, foreignRequire)
     }
 
     return js
@@ -357,20 +356,30 @@ function bundle(options, cache) {
 }
 
 // map of PS module names to their source path
-function psModuleMap(globs, cache) {
+function psModuleMap(options, cache) {
   if (cache.psModuleMap) return Promise.resolve(cache.psModuleMap)
 
+  const globs = [].concat(options.src).concat(options.ffi)
+
   return globby(globs).then(paths => {
     return Promise
       .props(paths.reduce((map, file) => {
         map[file] = fs.readFileAsync(file, 'utf8')
         return map
       }, {}))
-      .then(srcMap => {
-        cache.psModuleMap = Object.keys(srcMap).reduce((map, file) => {
-          const source = srcMap[file]
-          const psModuleName = match(psModuleRegex, source)
-          map[psModuleName] = path.resolve(file)
+      .then(fileMap => {
+        cache.psModuleMap = Object.keys(fileMap).reduce((map, file) => {
+          const source = fileMap[file]
+          const ext = path.extname(file)
+          const isPurs = ext.match(/purs$/i)
+          const moduleRegex = isPurs ? srcModuleRegex : ffiModuleRegex
+          const moduleName = match(moduleRegex, source)
+          map[moduleName] = map[moduleName] || {}
+          if (isPurs) {
+            map[moduleName].src = path.resolve(file)
+          } else {
+            map[moduleName].ffi = path.resolve(file)
+          }
           return map
         }, {})
         return cache.psModuleMap