Unfinished Lua sanitizer

lua.moon

:Stack, :in_array, :find_replace, :call = require './utils'
:concat = table

-- 1. Turn unnested `local =` declarations into `env.` mutations
-- 2. Track scoped variables (function params, for..in vars, explicit locals)
-- 3. Ensure all variable refs use `env.` when necessary
-- 4. Rename local vars named "env" or "res" to avoid conflicts

reserved_vars = {'res', 'env'}
scope_keywords = {'do', 'for', 'function', 'repeat', 'then'}
other_keywords = {
  'and', 'break', 'else', 'elseif', 'end', 'false', 'if', 'in',
  'nil', 'not', 'or', 'return', 'true', 'until', 'while',
}

-- Scope stack mutation.
push_env = => @env = setmetatable {}, __index: @env
pop_env = => @env = getmetatable(@env).__index

-- Reserved variables cannot be referenced or shadowed.
unreserved_var = (var) ->
  if in_array var, reserved_vars
    return var .. '$' .. 1

-- Prefix a variable with `env.` unless it exists locally.
prefix_var = (var, locals) ->
  return var if locals[var] ~= nil
  return 'env.' .. unreserved_var var

base_find = call ->
  tokens = Stack!

  add_token = (token, replace) ->
    tokens\push {token, replace}

  -- Possible variable name
  WORD = '[$_%a][$_%w]*'
  add_token WORD, (x, y, input) =>
    word = input\sub x, y

    -- Local refs are preserved.
    return if @env[word] ~= nil

    if in_array word, scope_keywords
      @scope\push word

      -- Add the loop vars to scope, and skip `do`.
      if word == 'for'
        -- TODO: Find/replace the `in` expression
        return

      -- Add the args to scope, and skip `do`.
      if word == 'function'
        return

      -- The others are easy peasy.
      return

    if word == 'local'

      -- Find the end of the line.
      line_end = input\find '\n', y + 2

      -- Reduce our search to the current line.
      input = input\sub y + 2, line_end

      -- Only top-level locals are replaced.
      top_level = @scope\get! == nil

      -- Look for an assignment operator.
      assign_idx = if top_level then input\find '='

      -- Nested locals and unassigned locals are preserved.
      if assign_idx == nil

        -- Add the declared variables to the scope.
        while true
          var_idx, var_end = input\find WORD
          if var_idx ~= nil
            var = input\sub var_idx, var_end
            @env[var] = true

        -- Skip to the next line.
        state.ch = line_end
        return

      -- Reduce our search to before assignment.
      input = input\sub 1, assign_idx - 1

      -- The replacement variables.
      vars = Stack!

      -- Assign variables to `env`
      while true
        var_idx, var_end = input\find WORD
        if var_idx ~= nil
          var = input\sub var_idx, var_end
          vars\push prefix_var var, @env
        else break

      -- The "local " part is removed.
      return concat(vars, ', '), x, assign_idx - 1

    scope = @scope\get!

    -- The `end` keyword pops the scope.
    if word == 'end'
      @scope\pop!
      return

    -- The `until` keyword pops the scope after a newline.
    if scope == 'repeat' and word == 'until'
      -- TODO: Find/replace the `until` expression
      @scope\pop!
      return

    -- All other keywords are skipped.
    unless in_array word, other_keywords

      if scope == '{'
        -- TODO: Look for "%s*=" immediately after the word.
        return

      -- TODO: Prefix the variable ref (unless local).
      return

  find_comment_end = (x, y) =>
    str = if y - x > 1 then '--%]%]' else '\n'
    x, y = input\find str, y + 1
    @ch = if y then y + 1 else nil

  -- Comment openers
  add_token '--%[%[', find_comment_end
  add_token '--', find_comment_end

  -- TODO: Look for escaped quotation mark.
  find_string_end = (x, y, input) =>
    str = if x == y then '%]%]' else input\sub x, y
    x, y = input\find str, y + 1
    @ch = if y then y + 1 else nil

  -- String openers
  add_token '["\']', find_string_end
  add_token '%[%[', find_string_end

  -- Table opener
  add_token '{', (x, y) =>
    @scope.push '{'
    @ch = y + 1

  -- Table closer
  add_token '}', (x, y) =>
    @scope.pop!
    @ch = y + 1

  return (input, ch) =>
    x, y = 0, 0
    for {token, replace} in *tokens
      a, b = input\find token, ch
      if a ~= nil
        x, y = a, b
        input = input\sub 1, a - 1
        @replace = replace

    if x > 0
      return x, y

-- Turn top-level locals into `env.` mutations
-- Ignore `local` declarations without assignment
-- Prefix `env.` to variable refs where necessary
sanitize = (code, env = {}) ->
  find_replace code,
    find: base_find
    scope: Stack!
    :unreserved_var
    :prefix_var
    :env

return {
  :sanitize
}

utils.moon

Stack = do
  __len: => @len
  __index:
    get: (i) => self[i or @len]
    push: (val) =>
      @len = @len + 1
      self[@len] = val
    pop: =>
      self[@len] = nil
      @len = @len - 1

setmetatable Stack,
  __call: -> setmetatable {len: 0}, Stack

in_array = (val, arr) ->
  i, x = 1, nil
  while true
    x = arr[i]
    return true if x == val
    return false if x == nil
    i += 1

-- The `replacer` is passed nil when no match is found
-- The `replacer` is responsible for incrementing `state.ch`
-- The search ends when `state.ch` equals nil or > input length
find_replace = (input, state) ->

  if type(state.find) ~= 'function'
    error '`state.find` must be a function'

  -- Where to begin the next search.
  state.ch = 1

  len = #input
  idx = 1 -- Where the unreplaced substring begins.
  out, n = {}, 0 -- The chunked output.

  while true
    -- Look for a substring.
    _x, _y = state\find input, state.ch

    -- Ask for a replacement.
    str, x, y = state\replace _x, _y, input

    -- A replacement is not guaranteed.
    if type(str) == 'string'

      -- A replacement cannot be used if no match was found.
      error 'Cannot replace nothing' if _x == nil

      -- The replacer may choose the replaced range.
      -- Otherwise, the matched range is replaced.
      x, y = _x, _y if x == nil

      -- Replacing characters before the matched substring is not allowed.
      error 'Cannot replace backwards' if x < _x

      -- Include any unreplaced characters.
      if x > idx
        n += 1
        out[n] = input\sub idx, x - 1

      -- Include the replacement string.
      n += 1
      out[n] = str

      -- The character after the replaced substring
      -- is the start of the next unreplaced substring.
      idx = y + 1

    -- Are we done searching?
    if state.ch == nil or state.ch > len

      -- Include any unreplaced characters.
      if idx <= len
        n += 1
        out[n] = str\sub idx, len

      -- Return the output string.
      return concat out, ''

return {
  :Stack
  :in_array
  :find_replace
  call: (fn) -> fn!
}