fcomplex.lua

--
--  Copyright (c) 2015, Facebook, Inc.
--  All rights reserved.
--
--  This source code is licensed under the BSD-style license found in the
--  LICENSE file in the root directory of this source tree. An additional grant
--  of patent rights can be found in the PATENTS file in the same directory.

-- Rudimentary complex number support, building on top of LuaJIT's already
-- existing support (1+2i is parsed correctly and produces a boxed float
-- _Complex value)
--
-- This provides basic mathematical operations and wrappers for the functions
-- in <complex.h>

local ffi = require('ffi')

local M = {}

local mt = {}
local complex

-- Ensure a number is of the complex type
local function tocomplex(a)
   if ffi.istype(complex, a) or ffi.istype('complex double', a) then
      return a
   elseif type(a) == 'number' then
      return complex(a, 0)
   else
      error('Invalid type, numeric type expected')
   end
end
M.tocomplex = tocomplex

-- Wrappers around <complex.h> functions
local one_operand = {
    'acos', 'acosh', 'arg', 'asin', 'asinh', 'atan', 'atanh', 'cos',
    'cosh', 'exp', 'log', 'sin', 'sinh', 'sqrt', 'tan', 'tanh',
}

local two_operand = {
    'pow',
}

local one_operand_return_real = { 'abs', 'arg' }

for _, fn in ipairs(one_operand) do
    local cname = 'c' .. fn .. 'f'
    ffi.cdef(string.format('float _Complex %s(float _Complex)', cname))
    M[fn] = function(a)
        return ffi.C[cname](tocomplex(a))
    end
end

for _, fn in ipairs(two_operand) do
    local cname = 'c' .. fn .. 'f'
    ffi.cdef(string.format(
        'float _Complex %s(float _Complex, float _Complex)', cname))
    M[fn] = function(a, b)
        return ffi.C[cname](tocomplex(a), tocomplex(b))
    end
end

for _, fn in ipairs(one_operand_return_real) do
    local cname = 'c' .. fn .. 'f'
    ffi.cdef(string.format('float %s(float _Complex)', cname))
    M[fn] = function(a)
       return ffi.C[cname](tocomplex(a))
    end
end



function M.conj(a)
    a = tocomplex(a)
    return complex(a.re, -a.im)
end

function mt.__add(a, b)
    a, b = tocomplex(a), tocomplex(b)
    return complex(a.re + b.re, a.im + b.im)
end

function mt.__sub(a, b)
    a, b = tocomplex(a), tocomplex(b)
    return complex(a.re - b.re, a.im - b.im)
end

function mt.__mul(a, b)
    a, b = tocomplex(a), tocomplex(b)
    return complex(a.re * b.re - a.im * b.im,
                   a.re * b.im + a.im * b.re)
end

function mt.__div(a, b)
    a, b = tocomplex(a), tocomplex(b)
    local d = b.re * b.re + b.im * b.im
    return complex((a.re * b.re + a.im * b.im) / d,
                   (a.im * b.re - a.re * b.im) / d)
end

mt.__pow = M.pow

function mt.__unm(a)
    a = tocomplex(a)
    return complex(-a.re, -a.im)
end

function mt.__eq(a, b)
   if b and a.re == b.re and a.im == b.im then
      return true
   else
      return false
   end
end

complex = ffi.metatype('float _Complex', mt)
M.type = complex

return M