GenericWASM.hs

{-# LANGUAGE QuasiQuotes #-}

module Futhark.CodeGen.Backends.GenericWASM
  ( GC.CParts (..),
    GC.asLibrary,
    GC.asExecutable,
    GC.asServer,
    EntryPointType,
    JSEntryPoint (..),
    JSRecordField (..),
    JSOpaqueType (..),
    emccExportNames,
    javascriptWrapper,
    extToString,
    opaqueToJS,
    runServer,
    libraryExports,
  )
where

import Data.List (intercalate)
import Data.Text qualified as T
import Futhark.CodeGen.Backends.GenericC qualified as GC
import Futhark.CodeGen.Backends.SimpleRep (isValidCName, opaqueName)
import Futhark.CodeGen.ImpCode.Sequential qualified as Imp
import Futhark.CodeGen.RTS.JavaScript
import Futhark.Util (nubOrd, showText, zEncodeText)
import Language.Futhark.Core (nameToText)
import Language.Futhark.Primitive
import NeatInterpolation (text)

extToString :: Imp.ExternalValue -> String
extToString (Imp.TransparentValue (Imp.ArrayValue vn _ pt s dimSize)) =
  concat (replicate (length dimSize) "[]") ++ extToString (Imp.TransparentValue (Imp.ScalarValue pt s vn))
extToString (Imp.TransparentValue (Imp.ScalarValue (FloatType Float16) _ _)) = "f16"
extToString (Imp.TransparentValue (Imp.ScalarValue (FloatType Float32) _ _)) = "f32"
extToString (Imp.TransparentValue (Imp.ScalarValue (FloatType Float64) _ _)) = "f64"
extToString (Imp.TransparentValue (Imp.ScalarValue (IntType Int8) Imp.Signed _)) = "i8"
extToString (Imp.TransparentValue (Imp.ScalarValue (IntType Int16) Imp.Signed _)) = "i16"
extToString (Imp.TransparentValue (Imp.ScalarValue (IntType Int32) Imp.Signed _)) = "i32"
extToString (Imp.TransparentValue (Imp.ScalarValue (IntType Int64) Imp.Signed _)) = "i64"
extToString (Imp.TransparentValue (Imp.ScalarValue (IntType Int8) Imp.Unsigned _)) = "u8"
extToString (Imp.TransparentValue (Imp.ScalarValue (IntType Int16) Imp.Unsigned _)) = "u16"
extToString (Imp.TransparentValue (Imp.ScalarValue (IntType Int32) Imp.Unsigned _)) = "u32"
extToString (Imp.TransparentValue (Imp.ScalarValue (IntType Int64) Imp.Unsigned _)) = "u64"
extToString (Imp.TransparentValue (Imp.ScalarValue Bool _ _)) = "bool"
extToString (Imp.TransparentValue (Imp.ScalarValue Unit _ _)) = error "extToString: Unit"
extToString (Imp.OpaqueValue oname _) = T.unpack $ opaqueName oname

type EntryPointType = String

data JSEntryPoint = JSEntryPoint
  { name :: String,
    parameters :: [EntryPointType],
    ret :: [EntryPointType]
  }

-- | A field in a JavaScript record opaque type.
data JSRecordField = JSRecordField
  { -- | Original field name (used in the server protocol).
    jsrfName :: String,
    -- | JavaScript type string (e.g. @"i32"@, @"[]f64"@, @"opaque_foo"@).
    jsrfType :: EntryPointType,
    -- | JavaScript method name on @FutharkContext@ (e.g. @"project_opaque_foo_x"@).
    jsrfProjectFn :: String
  }

-- | How an opaque type is represented for the JavaScript server.
data JSOpaqueType
  = -- | A record type with named, projectable fields.
    JSOpaqueRecord [JSRecordField]
  | -- | Any other opaque type (no project support).
    JSOpaqueOther

-- | Convert the opaque-type table from an ImpCode program into the
-- JavaScript-oriented representation used by the code generator.
opaqueToJS :: Imp.OpaqueTypes -> [(String, JSOpaqueType)]
opaqueToJS (Imp.OpaqueTypes types) = map convertOne types
  where
    convertOne (desc, Imp.OpaqueRecord fields) =
      (T.unpack (opaqueName desc), JSOpaqueRecord (map (convertField desc) fields))
    convertOne (desc, _) =
      (T.unpack (opaqueName desc), JSOpaqueOther)

    convertField desc (fname, etype) =
      JSRecordField
        { jsrfName = T.unpack (nameToText fname),
          jsrfType = entryTypeToJSString etype,
          jsrfProjectFn =
            "project_"
              ++ T.unpack (opaqueName desc)
              ++ "_"
              ++ T.unpack f'
        }
      where
        f'
          | isValidCName (opaqueName desc <> "_" <> nameToText fname) =
              nameToText fname
          | otherwise = zEncodeText (nameToText fname)

-- | Convert an 'Imp.EntryPointType' to its JavaScript string representation.
entryTypeToJSString :: Imp.EntryPointType -> String
entryTypeToJSString (Imp.TypeOpaque desc) = T.unpack $ opaqueName desc
entryTypeToJSString (Imp.TypeTransparent (Imp.ValueType sign (Imp.Rank rank) pt)) =
  concat (replicate rank "[]") ++ primToJSString sign pt
  where
    primToJSString _ (FloatType Float16) = "f16"
    primToJSString _ (FloatType Float32) = "f32"
    primToJSString _ (FloatType Float64) = "f64"
    primToJSString Imp.Signed (IntType Int8) = "i8"
    primToJSString Imp.Signed (IntType Int16) = "i16"
    primToJSString Imp.Signed (IntType Int32) = "i32"
    primToJSString Imp.Signed (IntType Int64) = "i64"
    primToJSString Imp.Unsigned (IntType Int8) = "u8"
    primToJSString Imp.Unsigned (IntType Int16) = "u16"
    primToJSString Imp.Unsigned (IntType Int32) = "u32"
    primToJSString Imp.Unsigned (IntType Int64) = "u64"
    primToJSString _ Bool = "bool"
    primToJSString _ Unit = error "entryTypeToJSString: Unit"

emccExportNames :: [JSEntryPoint] -> [(String, JSOpaqueType)] -> [String]
emccExportNames jses opaqueTypes =
  map (\jse -> "'_futhark_entry_" ++ T.unpack (GC.escapeName (T.pack (name jse))) ++ "'") jses
    ++ map (\arg -> "'" ++ gfn "new" arg ++ "'") arrays
    ++ map (\arg -> "'" ++ gfn "free" arg ++ "'") arrays
    ++ map (\arg -> "'" ++ gfn "shape" arg ++ "'") arrays
    ++ map (\arg -> "'" ++ gfn "values_raw" arg ++ "'") arrays
    ++ map (\arg -> "'" ++ gfn "values" arg ++ "'") arrays
    ++ map (\arg -> "'" ++ "_futhark_free_" ++ arg ++ "'") opaques
    ++ map (\rf -> "'_futhark_" ++ jsrfProjectFn rf ++ "'") allRecordFields
    ++ [ "_futhark_context_config_new",
         "_futhark_context_config_free",
         "_futhark_context_new",
         "_futhark_context_free",
         "_futhark_context_get_error"
       ]
  where
    -- Include array types from both entry points and record fields.
    arrays = nubOrd $ filter isArray (entryPointTypes ++ recordFieldTypes)
    -- Include opaque types from both entry points and record fields.
    opaques = nubOrd $ filter isOpaque (entryPointTypes ++ recordFieldTypes)
    entryPointTypes = concatMap (\jse -> parameters jse ++ ret jse) jses
    recordFieldTypes = [jsrfType rf | (_, JSOpaqueRecord fields) <- opaqueTypes, rf <- fields]
    gfn typ str = "_futhark_" ++ typ ++ "_" ++ baseType str ++ "_" ++ show (dim str) ++ "d"
    allRecordFields = [rf | (_, JSOpaqueRecord fields) <- opaqueTypes, rf <- fields]

javascriptWrapper :: [JSEntryPoint] -> [(String, JSOpaqueType)] -> T.Text
javascriptWrapper entryPoints opaqueTypes =
  T.unlines
    [ serverJs,
      valuesJs,
      wrapperclassesJs,
      classFutharkContext entryPoints opaqueTypes
    ]

classFutharkContext :: [JSEntryPoint] -> [(String, JSOpaqueType)] -> T.Text
classFutharkContext entryPoints opaqueTypes =
  T.unlines
    [ "class FutharkContext {",
      constructor entryPoints opaqueTypes,
      getFreeFun,
      getEntryPointsFun,
      getTypesFun,
      getErrorFun,
      T.unlines $ map toFutharkArray arrays,
      T.unlines $ concatMap (generateProjectMethods . snd) opaqueTypes,
      T.unlines $ map jsWrapEntryPoint entryPoints,
      "}",
      [text|
      async function newFutharkContext() {
        var wasm = await loadWASM();
        return new FutharkContext(wasm);
      }
      |]
    ]
  where
    -- Collect array types from entry points AND from record fields so the
    -- array constructors are always available when returning projected values.
    arrays = nubOrd $ filter isArray (entryPointTypes ++ recordFieldTypes)
    entryPointTypes = concatMap (\jse -> parameters jse ++ ret jse) entryPoints
    recordFieldTypes = [jsrfType rf | (_, JSOpaqueRecord fields) <- opaqueTypes, rf <- fields]

constructor :: [JSEntryPoint] -> [(String, JSOpaqueType)] -> T.Text
constructor jses opaqueTypes =
  [text|
  constructor(wasm, num_threads) {
    this.wasm = wasm;
    this.cfg = this.wasm._futhark_context_config_new();
    if (num_threads) this.wasm._futhark_context_config_set_num_threads(this.cfg, num_threads);
    this.ctx = this.wasm._futhark_context_new(this.cfg);
    this.entry_points = {
      ${entries}
    };
    this.types = {
      ${type_entries}
    };
  }
  |]
  where
    entries = T.intercalate "," $ map dicEntry jses
    type_entries = T.intercalate "," $ map dicTypeEntry opaqueTypes

getFreeFun :: T.Text
getFreeFun =
  [text|
  free() {
    this.wasm._futhark_context_free(this.ctx);
    this.wasm._futhark_context_config_free(this.cfg);
  }
  |]

getEntryPointsFun :: T.Text
getEntryPointsFun =
  [text|
  get_entry_points() {
    return this.entry_points;
  }
  |]

getTypesFun :: T.Text
getTypesFun =
  [text|
  get_types() {
    return this.types;
  }
  |]

getErrorFun :: T.Text
getErrorFun =
  [text|
  get_error() {
    var ptr = this.wasm._futhark_context_get_error(this.ctx);
    var len = this.wasm.HEAP8.subarray(ptr).indexOf(0);
    var str = String.fromCharCode(...this.wasm.HEAP8.subarray(ptr, ptr + len));
    this.wasm._free(ptr);
    return str;
  }
  |]

dicEntry :: JSEntryPoint -> T.Text
dicEntry jse =
  [text|
       "${ename}" : ["${fname}", ${params}, ${rets}]
  |]
  where
    fname = GC.escapeName $ T.pack $ name jse
    ename = T.pack $ name jse
    params = showText $ parameters jse
    rets = showText $ ret jse

-- | Generate the @types@ dictionary entry for a single opaque type.
dicTypeEntry :: (String, JSOpaqueType) -> T.Text
dicTypeEntry (tname, JSOpaqueRecord fields) =
  T.pack $
    show tname
      ++ ": [\"record\", ["
      ++ intercalate ", " (map fieldEntry fields)
      ++ "]]"
  where
    fieldEntry rf =
      "["
        ++ show (jsrfName rf)
        ++ ", "
        ++ show (jsrfType rf)
        ++ ", \""
        ++ jsrfProjectFn rf
        ++ "\"]"
dicTypeEntry (tname, JSOpaqueOther) =
  T.pack $ show tname ++ ": [\"opaque\"]"

-- | Generate @project_*@ wrapper methods for a record opaque type.
generateProjectMethods :: JSOpaqueType -> [T.Text]
generateProjectMethods (JSOpaqueRecord fields) = map generateProjectMethod fields
generateProjectMethods JSOpaqueOther = []

generateProjectMethod :: JSRecordField -> T.Text
generateProjectMethod rf =
  T.pack $
    unlines
      [ "  " ++ jsrfProjectFn rf ++ "(obj) {",
        "    var out = this.wasm._malloc(" ++ show (typeSize ftype) ++ ");",
        "    this.wasm._futhark_" ++ jsrfProjectFn rf ++ "(this.ctx, out, obj.ptr);",
        "    var result = " ++ readResult ++ ";",
        "    this.wasm._free(out);",
        "    return result;",
        "  }"
      ]
  where
    ftype = jsrfType rf
    readout = typeHeap ftype ++ "[out >> " ++ show (typeShift ftype) ++ "]"
    readResult
      | isArray ftype =
          "this.new_" ++ baseType ftype ++ "_" ++ show (dim ftype) ++ "d_from_ptr(" ++ readout ++ ")"
      | isOpaque ftype =
          "new FutharkOpaque(this, " ++ readout ++ ", this.wasm._futhark_free_" ++ ftype ++ ")"
      | ftype == "bool" = readout ++ "!==0"
      | otherwise = readout

jsWrapEntryPoint :: JSEntryPoint -> T.Text
jsWrapEntryPoint jse =
  [text|
  ${func_name}(${inparams}) {
    var out = [${outparams}].map(n => this.wasm._malloc(n));
    var to_free = [];
    var do_free = () => { out.forEach(this.wasm._free); to_free.forEach(f => f.free()); };
    ${paramsToPtr}
    if (this.wasm._futhark_entry_${func_name}(this.ctx, ...out, ${ins}) > 0) {
      do_free();
      throw this.get_error();
    }
    ${results}
    do_free();
    return ${res};
  }
  |]
  where
    func_name = GC.escapeName $ T.pack $ name jse

    alp = [0 .. length (parameters jse) - 1]
    inparams = T.pack $ intercalate ", " ["in" ++ show i | i <- alp]
    ins = T.pack $ intercalate ", " [maybeDerefence ("in" ++ show i) $ parameters jse !! i | i <- alp]
    paramsToPtr = T.pack $ unlines $ filter ("" /=) [arrayPointer ("in" ++ show i) $ parameters jse !! i | i <- alp]

    alr = [0 .. length (ret jse) - 1]
    outparams = T.pack $ intercalate ", " [show $ typeSize $ ret jse !! i | i <- alr]
    results = T.pack $ unlines [makeResult i $ ret jse !! i | i <- alr]
    res_array = intercalate ", " ["result" ++ show i | i <- alr]
    res = T.pack $ if length (ret jse) == 1 then "result0" else "[" ++ res_array ++ "]"

maybeDerefence :: String -> String -> String
maybeDerefence arg typ =
  if isScalar typ then arg else arg ++ ".ptr"

arrayPointer :: String -> String -> String
arrayPointer arg typ =
  if isArray typ
    then "  if (" ++ arg ++ " instanceof Array) { " ++ reassign ++ "; to_free.push(" ++ arg ++ "); }"
    else ""
  where
    reassign = arg ++ " = this.new_" ++ signature ++ "_from_jsarray(" ++ arg ++ ")"
    signature = baseType typ ++ "_" ++ show (dim typ) ++ "d"

makeResult :: Int -> String -> String
makeResult i typ =
  "  var result"
    ++ show i
    ++ " = "
    ++ if isArray typ
      then "this.new_" ++ signature ++ "_from_ptr(" ++ readout ++ ");"
      else
        if isOpaque typ
          then "new FutharkOpaque(this, " ++ readout ++ ", this.wasm._futhark_free_" ++ typ ++ ");"
          else readout ++ if typ == "bool" then "!==0;" else ";"
  where
    res = "out[" ++ show i ++ "]"
    readout = typeHeap typ ++ "[" ++ res ++ " >> " ++ show (typeShift typ) ++ "]"
    signature = baseType typ ++ "_" ++ show (dim typ) ++ "d"

baseType :: String -> String
baseType ('[' : ']' : end) = baseType end
baseType typ = typ

dim :: String -> Int
dim ('[' : ']' : end) = dim end + 1
dim _ = 0

isArray :: String -> Bool
isArray typ = take 2 typ == "[]"

isOpaque :: String -> Bool
isOpaque typ = take 6 typ == "opaque"

isScalar :: String -> Bool
isScalar typ = not (isArray typ || isOpaque typ)

typeSize :: String -> Integer
typeSize typ =
  case typ of
    "i8" -> 1
    "i16" -> 2
    "i32" -> 4
    "i64" -> 8
    "u8" -> 1
    "u16" -> 2
    "u32" -> 4
    "u64" -> 8
    "f16" -> 2
    "f32" -> 4
    "f64" -> 8
    "bool" -> 1
    _ -> 4

typeShift :: String -> Integer
typeShift typ =
  case typ of
    "i8" -> 0
    "i16" -> 1
    "i32" -> 2
    "i64" -> 3
    "u8" -> 0
    "u16" -> 1
    "u32" -> 2
    "u64" -> 3
    "f16" -> 1
    "f32" -> 2
    "f64" -> 3
    "bool" -> 0
    _ -> 2

typeHeap :: String -> String
typeHeap typ =
  case typ of
    "i8" -> "this.wasm.HEAP8"
    "i16" -> "this.wasm.HEAP16"
    "i32" -> "this.wasm.HEAP32"
    "i64" -> "this.wasm.HEAP64"
    "u8" -> "this.wasm.HEAPU8"
    "u16" -> "this.wasm.HEAPU16"
    "u32" -> "this.wasm.HEAPU32"
    "u64" -> "(new BigUint64Array(this.wasm.HEAP64.buffer))"
    "f16" -> "this.wasm.HEAPU16"
    "f32" -> "this.wasm.HEAPF32"
    "f64" -> "this.wasm.HEAPF64"
    "bool" -> "this.wasm.HEAP8"
    _ -> "this.wasm.HEAP32"

toFutharkArray :: String -> T.Text
toFutharkArray typ =
  [text|
  ${new}_from_jsarray(${arraynd_p}) {
    return this.${new}(${arraynd_flat_p}, ${arraynd_dims_p});
  }
  ${new}(array, ${dims}) {
    console.assert(array.length === ${dims_multiplied}, 'len=%s,dims=%s', array.length, [${dims}].toString());
      var copy = this.wasm._malloc(array.length << ${shift});
      ${heapType}.set(array, copy >> ${shift});
      var ptr = ${fnew}(this.ctx, copy, ${bigint_dims});
      this.wasm._free(copy);
      return this.${new}_from_ptr(ptr);
    }

    ${new}_from_ptr(ptr) {
      return new FutharkArray(this, ptr, ${args});
    }
    |]
  where
    d = dim typ
    ftype = baseType typ
    heap = typeHeap ftype
    signature = ftype ++ "_" ++ show d ++ "d"
    new = T.pack $ "new_" ++ signature
    fnew = T.pack $ "this.wasm._futhark_new_" ++ signature
    fshape = "this.wasm._futhark_shape_" ++ signature
    fvalues = "this.wasm._futhark_values_raw_" ++ signature
    ffree = "this.wasm._futhark_free_" ++ signature
    arraynd = "array" ++ show d ++ "d"
    shift = showText (typeShift ftype)
    heapType = T.pack heap
    arraynd_flat = if d > 1 then arraynd ++ ".flat()" else arraynd
    arraynd_dims = intercalate ", " [arraynd ++ mult i "[0]" ++ ".length" | i <- [0 .. d - 1]]
    dims = T.pack $ intercalate ", " ["d" ++ show i | i <- [0 .. d - 1]]
    dims_multiplied = T.pack $ intercalate "*" ["Number(d" ++ show i ++ ")" | i <- [0 .. d - 1]]
    bigint_dims = T.pack $ intercalate ", " ["BigInt(d" ++ show i ++ ")" | i <- [0 .. d - 1]]
    mult i s = concat $ replicate i s
    (arraynd_p, arraynd_flat_p, arraynd_dims_p) = (T.pack arraynd, T.pack arraynd_flat, T.pack arraynd_dims)
    args = T.pack $ intercalate ", " ["'" ++ ftype ++ "'", show d, heap, fshape, fvalues, ffree]

-- | Javascript code that can be appended to the generated module to
-- run a Futhark server instance on startup.
runServer :: T.Text
runServer =
  [text|
   Module.onRuntimeInitialized = () => {
     var context = new FutharkContext(Module);
     var server = new Server(context);
     server.run();
   }|]

-- | The names exported by the generated module.
libraryExports :: T.Text
libraryExports = "export {newFutharkContext, FutharkContext, FutharkArray, FutharkOpaque};"