madl_main.mad

--[=[
 o-----------------------------------------------------------------------------o
 |
 | MAD environment (sandbox)
 |
 | Methodical Accelerator Design - Copyright (c) 2016+
 | Support: http://cern.ch/mad  - mad at cern.ch
 | Authors: L. Deniau, laurent.deniau at cern.ch
 | Contrib: -
 |
 o-----------------------------------------------------------------------------o
 | You can redistribute this file and/or modify it under the terms of the GNU
 | General Public License GPLv3 (or later), as published by the Free Software
 | Foundation. This file is distributed in the hope that it will be useful, but
 | WITHOUT ANY WARRANTY OF ANY KIND. See http://gnu.org/licenses for details.
 o-----------------------------------------------------------------------------o

  Purpose:
  - Load in order all the modules in the MAD environment.

 o-----------------------------------------------------------------------------o
]=]

local M = {}

-- modules --------------------------------------------------------------------o

-- object model
local objmod = 'object' ! 'objalt'

-- list of modules to import in MAD (order matters!)
local modules = {
  -- utilities
  'typeid', 'constant', 'gmath', 'gfunc', 'gutil',
  -- maths
  'range', 'complex', 'matrix', 'mono', 'gtpsa', 'damap',
  -- physics
  'gphys', 'geomap', 'dynmap', 'symint', 'aper', 'synrad', -- 'symintc',
  -- objects
  objmod, 'beam', 'beta0', 'element', 'sequence', 'mtable',
  -- commands
  'command', 'survey', 'track', 'cofind', 'twiss', 'match', 'correct', 'plot',
  -- environments
  'madx',
  -- shared libs and processes
  'libmadx', 'pymad',
}

-- list of eXternal modules to import in MAD
local xmodules = {
  -- Debugger (Hook), Strict, Reflect, LPEG, LuaFun, LuaFileSystem, LuaUnit
  'dbg', 'strict', 'reflect', 'regex', 'lfun', 'lfs', 'utest',
}

-- list of global variables to cleanup after import
local globals = {
  'lpeg', 're', 'lfs',
}

-- globals --------------------------------------------------------------------o

_IEC60559 = false -- see https://ieeexplore.ieee.org/document/9091348

-- locals ---------------------------------------------------------------------o

local ffi    = require 'ffi'
local _C     = require 'madl_cmad'
local typeid = require 'madl_typeid' .typeid

local is_nil, is_boolean, is_number, is_string, is_mappable,
      is_callable, wprotect, wrestrict, set_concept, concept     in typeid
local not_iterable                                               in concept

-- override os.exit to close the Lua state and force the GC to call finalizers

local os_exit = os.exit
os.exit = \code,close -> os_exit(code, is_nil(close) and true or close)

-- implementation -------------------------------------------------------------o

local _hlp = {} -- hidden key.

M[_hlp] = {}  -- backup for help
M._C    = _C  -- C symbol access

-- env

local env = _M ; _M = nil

env.arg            =  arg
env.is_tty         = \ -> _C.mad_stdin_is_tty   ~= 0
env.is_interactive = \ -> _C.mad_is_interactive ~= 0

M.env = wprotect(setmetatable(env, {__tostring := "MAD.env"}))

-- option

local option = {

  strfmt   = "% -25s",  -- string  format
  intfmt   = "% -10d",  -- integer format (includes imatrix)
  numfmt   = "% -.10g", -- number  format (includes  matrix and complex)
  prteps   = 1e-14,     -- number  global eps for printing
  hdrwidth = 18,        -- mtable header width
  colwidth = 18,        -- mtable column width

  nocharge = false,     -- ignore particles charges
  dontstab = false,     -- don't stabilize   one-turn map
  dontsymp = false,     -- don't symplectify one-turn map
  ptcmodel = false,     -- use PTC model (quad, sext, ...)

  madxenv  = false,     -- inside MAD-X environment if true (see MADX:load)
  objmodel = objmod,    -- 'object' or 'objalt'

  debug    = nil,       -- proxy for mad_trace_level
  location = nil,       -- proxy for mad_trace_location
  fortid   = nil,       -- proxy for mad_trace_fortid
}

local coption = {
  debug    = "mad_trace_level",
  location = "mad_trace_location",
  fortid   = "mad_trace_fortid",
}

M.option = setmetatable(option, {
  __tostring = \-> "MAD.option",
  __index    = \_,k -> _C[assert(coption[k],
                            "invalid read access to object <MAD.option>")],
  __newindex = \_,k,v => _C[assert(coption[k],
                            "invalid write access to object <MAD.option>")] = v
                      end,
})

-- trace, warning

local ctrace = trace ; trace = nil
local cwarn  = warn  ; warn  = nil

M.warn  = \fmt,...     => cwarn(string.format(fmt, ...)) end
M.trace = \lvl,fmt,... => if lvl >= _C.mad_trace_level then
                            ctrace(lvl, string.format(fmt, ...))
                          end end

-- import

local loaded_mod = {}

local function loadreq (modnam)
  local status, mod = xpcall(require, debug.traceback, modnam)
  local err

  if not status then
    if not string.find(mod, "module '"..modnam.."' not found:") then
      error('<'..mod .. '> while loading module ' .. modnam)
    end
    mod, err = nil, mod
  end
  return mod, err
end

local function loadmod(modnam)
  local mod, hlp, err

  hlp, mod, err = true, loadreq('madl_'..modnam)    -- 1st attempt
  if not mod then
    hlp, mod, err = false, loadreq(modnam)          -- 2nd attempt
  end
  assert(mod, err)

  if loaded_mod[mod] then return nil end

  if hlp then
    hlp, err = loadreq('madh_'..modnam)                 -- 1st attempt for help
    if not hlp then hlp, err = loadreq('help.madh_'..modnam) end -- 2nd attempt
    assert(hlp, err)
    assert(not mod.__help, "help already registered")
    mod.__help = hlp
  end

  loaded_mod[mod] = mod
  return mod
end

function M.import (to, from, override_)
  if to ~= MAD and is_nil(override_) then
    to, from, override_ = MAD, to, from -- right shift
  end

  -- load module first
  if is_string(from) then
    from = loadmod(from)
    if is_nil(from) then return MAD end -- module already loaded
  end

  local override = override_ or false
  assert(to == MAD           , "invalid argument #1 (MAD expected)")
  assert(is_mappable(from)   , "invalid argument #2 (mappable expected)")
  assert(is_boolean(override), "invalid argument #3 (boolean expected)")

  -- collect and check first (make a copy in case of error)
  local obj = {}
  for k,v in pairs(from) do
    if is_string(k) and string.sub(k,1,2) ~= '__' then
      if M.option.debug >= 3 then
        io.write(string.format('import: %-25s %-25s\n',k,v))
      end
      assert(is_nil(MAD[k]) or MAD[k] == v or override == true,
             "cannot override definition for " .. tostring(k))
      obj[k] = v
    end
  end
  -- finalize the import
  for k,v in pairs(obj) do M[k] = v end
  -- record the help
  if is_mappable(from.__help) then
    M[_hlp][from.__help] = from.__help
  end
  return MAD
end

-- export

function M.export (from, what_, to_, override_)
  if from ~= MAD and is_nil(override_) then
    from, what_, to_, override_ = MAD, from, what_, to_ -- right shift
  end
  if (is_nil(what_) or is_boolean(what_)) and is_nil(override_) then
    what_, override_ = MAD, what_ -- right shift
  end
  if (is_nil(to_) or is_boolean(to_)) and is_nil(override_) then
    to_, override_ = _G, to_ -- right shift
  end
  local what     = is_string(what_) and {what_} or what_ or MAD
  local to       = to_       or _G
  local override = override_ or false
  assert(from == MAD         , "invalid argument #1 (MAD expected)")
  assert(is_mappable(what)   , "invalid argument #2 (mappable expected)")
  assert(is_mappable(to)     , "invalid argument #3 (mappable expected)")
  assert(is_boolean(override), "invalid argument #4 (boolean expected)")

  -- collect and check first (make a copy in case of error)
  local obj = {}
  for k,v in pairs(what) do
    if is_number(k) and is_string(v) then
      k, v = v, MAD[v]
    end
    if is_string(k) and string.sub(k,1,2) ~= '__' then
      if M.option.debug >= 3 then
        io.write(string.format('export: %-25s %-25s\n',k,v))
      end
      assert(is_nil(to[k]) or to[k] == v or override == true,
             "cannot override definition for " .. tostring(k))
      obj[k] = v
    end
  end

  -- finalize the export
  for k,v in pairs(obj) do to[k] = v end
  return MAD
end

-- help

local function help_grab (tt)
  local lst = {}
  for _,t in pairs(tt) do
  for k,v in pairs(t ) do
      assert(is_nil(lst[k]), "unexpected duplicated data for "..k)
      lst[k] = v
  end end
  return lst
end

local function help_fmt (tt, n)
  local len = 2 -- tab
  for i,v in ipairs(tt) do
    len = len + #v + 2 -- sep
    if len > n then tt[i] = '\n  '..v ; len = #v + 2 end
  end
  return table.concat(tt, ', ')
end

function M.help (from, pattern_)
  if from ~= MAD and is_nil(pattern_) then
    from, pattern_ = MAD, from  -- right shift
  end
  local pattern = pattern_ or "^([%w_]+)$"
  assert(from == MAD       , "invalid argument #1 (MAD expected)")
  assert(is_string(pattern), "invalid argument #2 (string expected)")

  local hlp, sel, key = help_grab(MAD[_hlp]), {}, nil
  for k,v in pairs(hlp) do
    assert(is_string(v), "invalid data for "..k.." (string expected)")
    local kl = string.lower(k)
    if k == pattern or kl == pattern then
      key = k
    elseif string.match(k, pattern) or string.match(kl, pattern) then
      sel[#sel+1] = k
    end
  end

  if #sel == 0 and is_nil(key) then
    io.write("  No help found for '" .. pattern .. "'.")
  elseif not is_nil(key) then
    io.write(hlp[key])
  end

  table.sort(sel)
  local str = help_fmt(sel, 80)
  if str ~= '' then
    io.write("\nRelated topics:\n  ", str, ".\n")
  end
  io.write("\n")
end

-- show

local show_type = {
  -- false
  table = nil, lightuserdata = nil, userdata = nil, cdata = nil,
  ['function'] = nil, thread = nil,
  -- true
  ['nil'] = true,  boolean = true, number = true, string = true,
}

local function show_cmp (ka, kb)
  if type(ka) == type(kb) then
    return ka < kb
  else
    return tostring(ka) < tostring(kb)
  end
end

function M.show (a, cls_)
  local is_object, is_cdata, get_metatable in MAD.typeid
  local luatostring                        in MAD.utility
      if is_object  (a) then io.write(':'..tostring(a), '\n')
  elseif is_mappable(a) then io.write(':'..luatostring(a), '\n')
  else io.write(':'..type(a)..': ', tostring(a), '\n')
  end

  local keys
  if is_object(a) then
    keys = a:get_varkeys(cls_)
  elseif is_mappable(a) or is_cdata(a) then
    if is_cdata(a) then
      local mt = get_metatable(a)
      if mt then a = mt end
    end
    keys = {}
    for k in pairs(a) do
      if is_string(k) and k:sub(1,2) ~= '__' then
        keys[#keys+1] = k
      end
    end
  end

  if keys then
    table.sort(keys, show_cmp)
    for i=1,#keys do
      local k,v = keys[i], a[keys[i]]
      local t = type(v)
      io.write( string.format("%-18s :", tostring(k)) )
      if show_type[t] then io.write(t, ': ') end
      io.write(tostring(v), '\n')
    end
  end
end

-- load environment -----------------------------------------------------------o

-- protect MAD
MAD = wprotect(setmetatable(M, {__tostring := "MAD"}))

-- import MAD as a self reference
MAD:import { MAD = MAD }

-- load eXternal modules
for _,m in ipairs(xmodules) do MAD:import(m) end

-- activate strict mode (see _G_env below)
MAD.strict({MADX=true, help=true, show=true})

-- load MAD modules
for _,m in ipairs(modules) do MAD:import(m) end

-- cleanup globals
for i,v in ipairs(globals) do _G[v] = nil end

-- useful as globals
help = MAD.help
show = MAD.show

-- check _G pollution ---------------------------------------------------------o

local _G_env = {
  -- Lua
  '_G', '_PROMPT', '_PROMPT2', '_VERSION',
  'arg', 'assert',
  'collectgarbage', 'coroutine',
  'debug', 'dofile',
  'error',
  'getmetatable',
  'io', 'ipairs',
  'load', 'loadfile', 'loadstring',
  'math', 'module',
  'next',
  'os',
  'package', 'pairs', 'pcall', 'print',
  'rawequal', 'rawget', 'rawlen', 'rawset', 'require',
  'select', 'setmetatable', 'string',
  'table', 'tonumber', 'tostring', 'type',
  'unpack',
  'xpcall',
  -- 5.1
  'gcinfo', 'getfenv', 'setfenv',
  -- 5.2
  'bit32',
  -- 5.3
  'utf8',
  -- LJ 2.1
  'bit', 'jit', 'newproxy',
  -- MAD
  'MAD', 'MADX', '_IEC60559', 'help', 'show',
}
for i,v in ipairs(_G_env) do _G_env[v] = i end

local function check_G (extra)
  extra = extra or {}
  for k,v in pairs(_G) do
    if not (_G_env[k] or extra[k]) then
      io.write("unexpected global variable: ",
               tostring(k), ' -> ', tostring(v), '\n')
    end
  end
end

-- check MAD pollution --------------------------------------------------------o

local _M_env = {
  '_C', 'MAD', 'MADX', 'checkglobal',
  'aperture',
  'beam', 'beta0',
  'cdamap', 'cmatrix', 'cofind', 'command', 'complex', 'constant', 'correct',
  'ctpsa', 'cvector', 'cvname',
  'damap', 'dbg', 'dynmap',
  'element', 'env', 'export',
  'filesys', 'lfun',
  'geomap', 'gfunc', 'gmath', 'gphys', 'gplot', 'gtpsad', 'gtpsad_del',
  'help',
  'imatrix', 'import', 'ivector',
  'libmadx', 'linspace', 'logrange', 'logspace',
  'match', 'matrix', 'monomial', 'mtable',
  'nlogrange', 'nrange',
  'object', 'option',
  'plot', 'pymad',
  'range', 'reflect', 'regex',
  'sequence', 'show', 'strict', 'survey', 'symint', 'symintc', 'synrad',
  'toboolean', 'tocomplex', 'torange', 'tostring', 'totable',
  'tpsa', 'trace', 'track', 'twiss', 'typeid',
  'utest', 'utility',
  'vector',
  'warn',
  [_hlp] = true,
}
for i,v in ipairs(_M_env) do _M_env[v] = i end

local function check_M ()
  for k,v in pairs(MAD) do
    if not _M_env[k] then
      io.write("unexpected MAD variable: ",
               tostring(k), ' -> ', tostring(v), '\n')
    end
  end
end


-- pollution checks -----------------------------------------------------------o

M.checkglobal = \extra => check_G(extra) ; check_M() end

-- execute after loading MAD env
M.checkglobal()

-- end ------------------------------------------------------------------------o