--- Core array module.

-- Standard library imports --

local error = error

local getmetatable = getmetatable

local rawequal = rawequal

local setmetatable = setmetatable

local tostring = tostring

local traceback = debug and debug.traceback

local type = type

-- Modules --

local af = require("arrayfire_lib")

-- Cached module references --

local _AddToCurrentEnvironment_

local _Call_

local _CallWrap_

local _CheckError_

local _GetConstant_

local _GetFNSD_

local _IsArray_

local _IsConstant_

local _ToArray_

local _WrapArray_

-- Exports --

local M = {}

-- --

local ArrayMethodsAndMetatable = {}

local MetaValue = {}

local Constants = setmetatable({}, { __mode = "k" })

--

local function ErrorOut (what)

if traceback then

print(traceback())

end

error(what)

end

--

local function CallFromName_Checked (name, ...)

local func = af[name]

if type(func) ~= "function" then

if type(name) ~= "string" then

ErrorOut("Expected string name, got: " .. tostring(name))

else

ErrorOut(name .. " is not a function")

end

end

Name = name

return _CheckError_(func(...))

end

--

local function CallFromName_Unchecked (name, ...)

Name = name

return _CheckError_(af[name](...))

end

-- --

local CallFromName

--- DOCME

-- @string name

-- @param ... Arguments to function.

-- @return Any non-error return values.

function M.Call (name, ...)

return CallFromName(name, ...)

end

--

local function WrapAndReturn (arr, ...)

return _WrapArray_(arr), ...

end

--- DOCME

-- @string name

-- @param ... Arguments to function.

-- @treturn LuaArray X

-- @return Any additional return values.

function M.CallWrap (name, ...)

return WrapAndReturn(CallFromName(name, ...))

end

-- --

local SUCCESS = af.AF_SUCCESS

--- DOCME

-- @tparam af_err err

-- @param ...

-- @return ...

function M.CheckError (err, ...)

if err ~= SUCCESS then

local name = Name or ""

Name = nil

error(("%s: %i"):format(name, err))

end

return ...

end

--- DOCME

function M.CheckNames (check)

CallFromName = check and CallFromName_Checked or CallFromName_Unchecked

end

-- --

local Dim = {}

--- DOCME

-- @tparam af_array ha

-- @int dim

-- @treturn int FNSD

function M.GetFNSD (ha, dim)

if dim < 0 then

local ndims = _Call_("af_get_numdims", ha)

Dim[1], Dim[2], Dim[3], Dim[4] = _Call_("af_get_dims", ha)

for i = 1, 4 do

if Dim[i] > 1 then

return i - 1

end

end

return 0

else

return dim

end

end

--- DOCME

-- @tparam Constant k

-- @return R

function M.GetConstant (k)

return Constants[k] and k[1]

end

--- DOCME

-- @tparam LuaArray arr

-- @treturn ?|af_array|nil X

function M.GetHandle (arr)

-- TODO: If proxy, add reference?

return arr.m_handle

end

-- --

local Lib

--- DOCME

function M.GetLib ()

Lib = Lib or require("arrayfire")

return Lib

end

--- DOCME

-- @tparam function func

-- @tparam LuaArray arr

-- @number[opt=-1] dim

-- @string[opt] how

-- @return Results of _func_.

function M.HandleDim (func, arr, dim, how)

local harr = arr:get()

return (how ~= "no_wrap" and _CallWrap_ or _Call_)(func, harr, _GetFNSD_(harr, dim or -1))

end

--- DOCME

-- @param item

-- @treturn boolean B

function M.IsArray (item)

return rawequal(getmetatable(item), MetaValue) and not Constants[item]

end

--- DOCME

-- @param item

-- @treturn boolean B

function M.IsConstant (item)

return not not Constants[item] -- metatable redundant; coerce to false if missing

end

-- TODO: IsProxy(), MakeProxy()...

--- DOCME

-- @tparam LuaArray arr

-- @tparam ?|af_array|nil handle

function M.SetHandle (arr, handle)

-- TODO: disable for proxies

local cur = arr.m_handle

if cur ~= nil then

_Call_("af_release_array", cur)

end

arr.m_handle = handle

end

-- --

local Args = {}

--- DOCME

-- @param value

-- @tparam LuaArray other

-- @treturn LuaArray A

function M.ToArray (value, other)

if _IsConstant_(value) then

value = _GetConstant_(value)

end

local btype, hother = type(value), other:get()

local ndims = _Call_("af_get_numdims", hother)

Args[1], Args[2], Args[3], Args[4] = _Call_("af_get_dims", hother)

if btype == "table" then

-- Complex...

elseif btype == "number" then

-- Argh... detect range, integer-ness, etc?

return _Call_("af_constant", value, ndims, Args, af.f32)

end

end

--- DOCME

-- @string ret_type

-- @param real

-- @param imag

-- @treturn RetType RT

function M.ToType (ret_type, real, imag)

if rtype == "c32" or rtype == "c64" then

return { real = real, imag = imag }

else

return real -- TODO: Improve this!

end

end

--- DOCME

-- @string name

-- @param a

-- @param b

-- @param ...

-- @treturn LuaArray X

function M.TwoArrays (name, a, b, ...)

local atemp, btemp, ha, hb

if not _IsArray_(a) then

hb, ha, atemp = b:get(), _ToArray_(a, b), true

elseif not _IsArray_(b) then

ha, hb, btemp = a:get(), _ToArray_(b, a), true

else

ha, hb = a:get(), b:get()

end

--

local err, arr = af[name](ha, hb, ...)

if atemp or btemp then

_Call_("af_release_array", atemp and ha or hb)

end

Name = name

_CheckError_(err)

return _WrapArray_(arr)

end

--- DOCME

-- @tparam af_array arr

-- @treturn LuaArray X

function M.WrapArray (arr)

local wrapped = setmetatable({ m_handle = arr }, ArrayMethodsAndMetatable)

_AddToCurrentEnvironment_("array", wrapped)

return wrapped

end

--- DOCME

-- @param k constant

-- @treturn Constant Y

function M.WrapConstant (k)

k = { k }

Constants[k] = true

return setmetatable(k, ArrayMethodsAndMetatable) -- allow comparison operators

end

--

for _, v in ipairs{

"impl.ephemeral",

"impl.operators",

"impl.seq", -- depends on ephemeral

"impl.index", -- depends on ephemeral, seq

"methods.methods"

} do

require(v).Add(M, ArrayMethodsAndMetatable)

end

ArrayMethodsAndMetatable.__index = ArrayMethodsAndMetatable

ArrayMethodsAndMetatable.__metatable = MetaValue

-- Register array environment type.

M.RegisterEnvironmentCleanup("array", function(arr)

local ha = arr:get()

arr.m_handle = nil -- set() can error out

return not ha or af.af_release_array(ha) == SUCCESS

-- TODO: pooling?

end, "Errors releasing %i arrays")

-- TODO: Register "array_proxy"?

-- By default, check valid names.

M.CheckNames(true)

-- Cache module members.

_AddToCurrentEnvironment_ = M.AddToCurrentEnvironment

_Call_ = M.Call

_CallWrap_ = M.CallWrap

_CheckError_ = M.CheckError

_GetConstant_ = M.GetConstant

_GetFNSD_ = M.GetFNSD

_IsArray_ = M.IsArray

_IsConstant_ = M.IsConstant

_ToArray_ = M.ToArray

_WrapArray_ = M.WrapArray

-- Export the module.

return M