luau-vm 0.732.0

use crate::VmResult;
use crate::native::{NativeCallContext, NativeCallResult};
use crate::state::ThreadState;
use crate::thread::stack::RawStackAccess;
use crate::thread::{LUAI_MAX_NATIVE_CALLS, LuaStringBuilder, LuaStringBuilderStorage, Thread};
use crate::types::{LUA_TFUNCTION, LUA_TNUMBER, LUA_TSTRING, LUA_TTABLE};
use crate::vm::VmOperations;

use super::{L_ESC, digit, pos_relat};

const CAP_UNFINISHED: isize = -1;
const CAP_POSITION: isize = -2;
const LUA_MAX_CAPTURES: usize = 32;
const SPECIALS: &[u8] = b"^$*+?.([%-";

#[derive(Clone, Copy)]
struct Capture {
    init: usize,
    len: isize,
}

struct MatchState<'thread, 'a> {
    thread: &'thread Thread,
    match_depth: i32,
    src: &'a [u8],
    pat: &'a [u8],
    level: usize,
    capture: [Capture; LUA_MAX_CAPTURES],
}

/// `ms_byte`
fn match_state_byte(state: &MatchState<'_, '_>, index: usize) -> u8 {
    state.src.get(index).copied().unwrap_or(0)
}

/// `check_capture`
fn check_capture(state: &MatchState<'_, '_>, level: u8) -> VmResult<usize> {
    let level_index = i32::from(level) - i32::from(b'1');
    let index = level_index.max(0) as usize;
    if level < b'1' || index >= state.level || state.capture[index].len == CAP_UNFINISHED {
        return unsafe {
            crate::error!(state.thread, "invalid capture index %%%d", level_index + 1)
        }
        .map_err(Into::into);
    }
    Ok(index)
}

/// `capture_to_close`
fn capture_to_close(state: &MatchState<'_, '_>) -> VmResult<usize> {
    for level in (0..state.level).rev() {
        if state.capture[level].len == CAP_UNFINISHED {
            return Ok(level);
        }
    }

    unsafe { crate::error!(state.thread, "invalid pattern capture") }.map_err(Into::into)
}

/// `classend`
fn class_end(state: &MatchState<'_, '_>, mut index: usize) -> VmResult<usize> {
    let byte = state.pat[index];
    index += 1;

    match byte {
        L_ESC => {
            if index == state.pat.len() {
                return unsafe { crate::error!(state.thread, "malformed pattern (ends with '%')") }
                    .map_err(Into::into);
            }
            Ok(index + 1)
        }
        b'[' => {
            if state.pat.get(index) == Some(&b'^') {
                index += 1;
            }

            loop {
                if index == state.pat.len() {
                    return unsafe {
                        crate::error!(state.thread, "malformed pattern (missing ']')")
                    }
                    .map_err(Into::into);
                }

                let current = state.pat[index];
                index += 1;

                if current == L_ESC && index < state.pat.len() {
                    index += 1;
                }

                if state.pat.get(index) == Some(&b']') {
                    break;
                }
            }

            Ok(index + 1)
        }
        _ => Ok(index),
    }
}

/// `match_class`
fn match_class(byte: u8, class: u8) -> bool {
    let result = match class.to_ascii_lowercase() {
        b'a' => byte.is_ascii_alphabetic(),
        b'c' => byte.is_ascii_control(),
        b'd' => byte.is_ascii_digit(),
        b'g' => byte.is_ascii_graphic(),
        b'l' => byte.is_ascii_lowercase(),
        b'p' => byte.is_ascii_punctuation(),
        b's' => byte.is_ascii_whitespace(),
        b'u' => byte.is_ascii_uppercase(),
        b'w' => byte.is_ascii_alphanumeric(),
        b'x' => byte.is_ascii_hexdigit(),
        b'z' => byte == 0,
        _ => return class == byte,
    };

    if class.is_ascii_lowercase() {
        result
    } else {
        !result
    }
}

/// `matchbracketclass`
fn match_bracket_class(byte: u8, pattern: &[u8]) -> bool {
    let mut matches = true;
    let mut index = 1usize;

    if pattern.get(index) == Some(&b'^') {
        matches = false;
        index += 1;
    }

    while index < pattern.len() - 1 {
        let current = pattern[index];
        if current == L_ESC {
            index += 1;
            if match_class(byte, pattern[index]) {
                return matches;
            }
        } else if index + 2 < pattern.len().saturating_sub(1) && pattern[index + 1] == b'-' {
            index += 2;
            if pattern[index - 2] <= byte && byte <= pattern[index] {
                return matches;
            }
        } else if current == byte {
            return matches;
        }

        index += 1;
    }

    !matches
}

/// `singlematch`
fn single_match(
    state: &MatchState<'_, '_>,
    src_index: usize,
    pat_index: usize,
    pat_end: usize,
) -> bool {
    if src_index >= state.src.len() {
        return false;
    }

    let byte = state.src[src_index];
    match state.pat[pat_index] {
        b'.' => true,
        L_ESC => match_class(byte, state.pat[pat_index + 1]),
        b'[' => match_bracket_class(byte, &state.pat[pat_index..pat_end]),
        pattern_byte => pattern_byte == byte,
    }
}

/// `matchbalance`
fn match_balance(
    state: &MatchState<'_, '_>,
    src_index: usize,
    pat_index: usize,
) -> VmResult<Option<usize>> {
    if pat_index >= state.pat.len() - 1 {
        return unsafe {
            crate::error!(
                state.thread,
                "malformed pattern (missing arguments to '%b')"
            )
        }
        .map_err(Into::into);
    }

    if match_state_byte(state, src_index) != state.pat[pat_index] {
        return Ok(None);
    }

    let begin = state.pat[pat_index];
    let end = state.pat[pat_index + 1];
    let mut depth = 1;
    let mut index = src_index + 1;

    while index < state.src.len() {
        let current = state.src[index];
        if current == end {
            depth -= 1;
            if depth == 0 {
                return Ok(Some(index + 1));
            }
        } else if current == begin {
            depth += 1;
        }
        index += 1;
    }

    Ok(None)
}

/// `max_expand`
fn max_expand(
    state: &mut MatchState<'_, '_>,
    src_index: usize,
    pat_index: usize,
    pat_end: usize,
) -> VmResult<Option<usize>> {
    let mut count = 0usize;
    while single_match(state, src_index + count, pat_index, pat_end) {
        count += 1;
    }

    loop {
        if let Some(result) = match_pattern(state, src_index + count, pat_end + 1)? {
            return Ok(Some(result));
        }

        if count == 0 {
            break;
        }
        count -= 1;
    }

    Ok(None)
}

/// `min_expand`
fn min_expand(
    state: &mut MatchState<'_, '_>,
    mut src_index: usize,
    pat_index: usize,
    pat_end: usize,
) -> VmResult<Option<usize>> {
    loop {
        if let Some(result) = match_pattern(state, src_index, pat_end + 1)? {
            return Ok(Some(result));
        } else if single_match(state, src_index, pat_index, pat_end) {
            src_index += 1;
        } else {
            return Ok(None);
        }
    }
}

/// `start_capture`
fn start_capture(
    state: &mut MatchState<'_, '_>,
    src_index: usize,
    pat_index: usize,
    kind: isize,
) -> VmResult<Option<usize>> {
    let level = state.level;
    if level >= LUA_MAX_CAPTURES {
        return unsafe { crate::error!(state.thread, "too many captures") }.map_err(Into::into);
    }

    state.capture[level] = Capture {
        init: src_index,
        len: kind,
    };
    state.level = level + 1;

    let result = match_pattern(state, src_index, pat_index)?;
    if result.is_none() {
        state.level -= 1;
    }

    Ok(result)
}

/// `end_capture`
fn end_capture(
    state: &mut MatchState<'_, '_>,
    src_index: usize,
    pat_index: usize,
) -> VmResult<Option<usize>> {
    let level = capture_to_close(state)?;
    state.capture[level].len = src_index as isize - state.capture[level].init as isize;

    let result = match_pattern(state, src_index, pat_index)?;
    if result.is_none() {
        state.capture[level].len = CAP_UNFINISHED;
    }

    Ok(result)
}

/// `match_capture`
fn match_capture(
    state: &MatchState<'_, '_>,
    src_index: usize,
    level: u8,
) -> VmResult<Option<usize>> {
    let level = check_capture(state, level)?;
    let len = state.capture[level].len as usize;

    if state.src.len().saturating_sub(src_index) >= len
        && state.src[state.capture[level].init..state.capture[level].init + len]
            == state.src[src_index..src_index + len]
    {
        Ok(Some(src_index + len))
    } else {
        Ok(None)
    }
}

/// `match`
fn match_pattern(
    state: &mut MatchState<'_, '_>,
    mut src_index: usize,
    mut pat_index: usize,
) -> VmResult<Option<usize>> {
    if state.match_depth == 0 {
        return unsafe { crate::error!(state.thread, "pattern too complex") }.map_err(Into::into);
    }
    state.match_depth -= 1;

    let global = unsafe { state.thread.global() };
    let interrupt = global.take_pattern_interrupt_callback();
    if let Some(interrupt) = interrupt {
        unsafe {
            state.thread.increment_native_call_depth();
            let result = interrupt(state.thread);
            state.thread.decrement_native_call_depth();
            result?;
        }
    }

    while pat_index < state.pat.len() {
        let result = match state.pat[pat_index] {
            b'(' => {
                if state.pat.get(pat_index + 1) == Some(&b')') {
                    start_capture(state, src_index, pat_index + 2, CAP_POSITION)?
                } else {
                    start_capture(state, src_index, pat_index + 1, CAP_UNFINISHED)?
                }
            }
            b')' => end_capture(state, src_index, pat_index + 1)?,
            b'$' if pat_index + 1 == state.pat.len() => {
                if src_index == state.src.len() {
                    Some(src_index)
                } else {
                    None
                }
            }
            L_ESC => match state.pat.get(pat_index + 1).copied().unwrap_or_default() {
                b'b' => {
                    if let Some(next) = match_balance(state, src_index, pat_index + 2)? {
                        src_index = next;
                        pat_index += 4;
                        continue;
                    }
                    None
                }
                b'f' => {
                    pat_index += 2;
                    if state.pat.get(pat_index) != Some(&b'[') {
                        return unsafe {
                            crate::error!(state.thread, "missing '[' after '%f' in pattern")
                        }
                        .map_err(Into::into);
                    }

                    let pat_end = class_end(state, pat_index)?;
                    let previous = if src_index == 0 {
                        0
                    } else {
                        state.src[src_index - 1]
                    };

                    if !match_bracket_class(previous, &state.pat[pat_index..pat_end])
                        && match_bracket_class(
                            match_state_byte(state, src_index),
                            &state.pat[pat_index..pat_end],
                        )
                    {
                        pat_index = pat_end;
                        continue;
                    } else {
                        None
                    }
                }
                b'0'..=b'9' => {
                    if let Some(next) = match_capture(state, src_index, state.pat[pat_index + 1])? {
                        src_index = next;
                        pat_index += 2;
                        continue;
                    }
                    None
                }
                _ => {
                    let pat_end = class_end(state, pat_index)?;
                    if !single_match(state, src_index, pat_index, pat_end) {
                        match state.pat.get(pat_end).copied() {
                            Some(b'*' | b'?' | b'-') => {
                                pat_index = pat_end + 1;
                                continue;
                            }
                            _ => None,
                        }
                    } else {
                        match state.pat.get(pat_end).copied() {
                            Some(b'?') => {
                                if let Some(result) =
                                    match_pattern(state, src_index + 1, pat_end + 1)?
                                {
                                    Some(result)
                                } else {
                                    pat_index = pat_end + 1;
                                    continue;
                                }
                            }
                            Some(b'+') => max_expand(state, src_index + 1, pat_index, pat_end)?,
                            Some(b'*') => max_expand(state, src_index, pat_index, pat_end)?,
                            Some(b'-') => min_expand(state, src_index, pat_index, pat_end)?,
                            _ => {
                                src_index += 1;
                                pat_index = pat_end;
                                continue;
                            }
                        }
                    }
                }
            },
            _ => {
                let pat_end = class_end(state, pat_index)?;
                if !single_match(state, src_index, pat_index, pat_end) {
                    match state.pat.get(pat_end).copied() {
                        Some(b'*' | b'?' | b'-') => {
                            pat_index = pat_end + 1;
                            continue;
                        }
                        _ => None,
                    }
                } else {
                    match state.pat.get(pat_end).copied() {
                        Some(b'?') => {
                            if let Some(result) = match_pattern(state, src_index + 1, pat_end + 1)?
                            {
                                Some(result)
                            } else {
                                pat_index = pat_end + 1;
                                continue;
                            }
                        }
                        Some(b'+') => max_expand(state, src_index + 1, pat_index, pat_end)?,
                        Some(b'*') => max_expand(state, src_index, pat_index, pat_end)?,
                        Some(b'-') => min_expand(state, src_index, pat_index, pat_end)?,
                        _ => {
                            src_index += 1;
                            pat_index = pat_end;
                            continue;
                        }
                    }
                }
            }
        };

        state.match_depth += 1;
        return Ok(result);
    }

    state.match_depth += 1;
    Ok(Some(src_index))
}

/// `lmemfind`
fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option<usize> {
    if needle.is_empty() {
        return Some(0);
    }
    if needle.len() > haystack.len() {
        return None;
    }

    haystack
        .windows(needle.len())
        .position(|window| window == needle)
}

/// `push_onecapture`
fn push_one_capture(
    state: &MatchState<'_, '_>,
    index: usize,
    start: usize,
    end: usize,
) -> VmResult {
    if index >= state.level {
        if index == 0 {
            unsafe { state.thread.push_string(&state.src[start..end])? };
        } else {
            return unsafe { crate::error!(state.thread, "invalid capture index") }
                .map_err(Into::into);
        }
    } else {
        let len = state.capture[index].len;
        if len == CAP_UNFINISHED {
            return unsafe { crate::error!(state.thread, "unfinished capture") }
                .map_err(Into::into);
        }
        if len == CAP_POSITION {
            unsafe {
                state
                    .thread
                    .push_integer(state.capture[index].init as i32 + 1)?
            };
        } else {
            let init = state.capture[index].init;
            unsafe {
                state
                    .thread
                    .push_string(&state.src[init..init + len as usize])?
            };
        }
    }
    Ok(())
}

/// `push_captures`
fn push_captures(state: &MatchState<'_, '_>, matched: Option<(usize, usize)>) -> VmResult<i32> {
    let count = if state.level == 0 && matched.is_some() {
        1
    } else {
        state.level as i32
    };

    unsafe {
        state
            .thread
            .lua_check_stack(count, Some("too many captures"))?
    };
    for index in 0..count as usize {
        let (start, end) = matched.unwrap_or((0, 0));
        push_one_capture(state, index, start, end)?;
    }

    Ok(count)
}

/// `nospecials`
fn has_no_specials(pattern: &[u8]) -> bool {
    !pattern.iter().any(|byte| SPECIALS.contains(byte))
}

/// `prepstate`
fn new_match_state<'thread, 'a>(
    thread: &'thread Thread,
    src: &'a [u8],
    pat: &'a [u8],
) -> MatchState<'thread, 'a> {
    MatchState {
        thread,
        match_depth: LUAI_MAX_NATIVE_CALLS as i32,
        src,
        pat,
        level: 0,
        capture: [Capture {
            init: 0,
            len: CAP_UNFINISHED,
        }; LUA_MAX_CAPTURES],
    }
}

/// `reprepstate`
fn reset_match_state(state: &mut MatchState<'_, '_>) {
    state.level = 0;
    debug_assert_eq!(state.match_depth, LUAI_MAX_NATIVE_CALLS as i32);
}

/// `str_find_aux`
fn string_find_aux(thread: &Thread, find: bool) -> NativeCallResult {
    let src = unsafe { thread.check_string(1)? };
    let pat = unsafe { thread.check_string(2)? };
    let mut init = pos_relat(unsafe { thread.opt_integer(3, 1)? }, src.len());

    if init < 1 {
        init = 1;
    } else if init > src.len() as i32 + 1 {
        unsafe { thread.push_nil()? };
        return Ok(1);
    }

    if find && (unsafe { thread.to_boolean(4) } != 0 || has_no_specials(pat)) {
        if let Some(found) = find_subslice(&src[(init - 1) as usize..], pat) {
            let start = found + (init - 1) as usize;
            unsafe {
                thread.push_integer(start as i32 + 1)?;
                thread.push_integer((start + pat.len()) as i32)?;
            }
            return Ok(2);
        }
    } else {
        let mut pattern = pat;
        let mut start = (init - 1) as usize;
        let anchor = pattern.first() == Some(&b'^');
        if anchor {
            pattern = &pattern[1..];
        }

        let mut state = new_match_state(thread, src, pattern);
        loop {
            reset_match_state(&mut state);
            if let Some(result) = match_pattern(&mut state, start, 0)? {
                if find {
                    unsafe {
                        thread.push_integer(start as i32 + 1)?;
                        thread.push_integer(result as i32)?;
                    }
                    return Ok((push_captures(&state, None)? + 2) as usize);
                }

                return Ok(push_captures(&state, Some((start, result)))? as usize);
            }

            if start >= src.len() || anchor {
                break;
            }
            start += 1;
        }
    }

    unsafe { thread.push_nil()? };
    Ok(1)
}

/// `str_find`
pub(super) fn string_find(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    string_find_aux(thread, true)
}

/// `str_match`
pub(super) fn string_match(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    string_find_aux(thread, false)
}

/// `gmatch_aux`
fn string_gmatch_aux(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let src = thread.check_string(crate::thread::upvalue_index(1))?;
        let pat = thread.check_string(crate::thread::upvalue_index(2))?;
        let mut state = new_match_state(thread, src, pat);
        let mut src_index = thread.opt_integer(crate::thread::upvalue_index(3), 0)? as usize;

        while src_index <= state.src.len() {
            reset_match_state(&mut state);
            if let Some(end) = match_pattern(&mut state, src_index, 0)? {
                let mut new_start = end as i32;
                if end == src_index {
                    new_start += 1;
                }

                thread.push_integer(new_start)?;
                thread.replace(crate::thread::upvalue_index(3));

                return Ok(push_captures(&state, Some((src_index, end)))? as usize);
            }

            src_index += 1;
        }

        Ok(0)
    }
}

/// `gmatch`
pub(super) fn string_gmatch(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let _ = thread.check_string(1)?;
        let _ = thread.check_string(2)?;
        thread.set_top(2)?;
        thread.push_integer(0)?;
        thread.push_native_closure_k(string_gmatch_aux, None, 3, None)?;
    }
    Ok(1)
}

/// `add_s`
unsafe fn add_s(
    state: &MatchState<'_, '_>,
    buffer: &mut LuaStringBuilder<'_, '_>,
    start: usize,
    end: usize,
) -> VmResult {
    unsafe {
        let replacement = state.thread.check_string(3)?;
        let _ = buffer.reserve(replacement.len())?;

        let mut index = 0usize;
        while index < replacement.len() {
            if replacement[index] != L_ESC {
                buffer.push_byte(replacement[index])?;
                index += 1;
                continue;
            }

            index += 1;
            if index == replacement.len() || !digit(replacement[index]) {
                if index == replacement.len() || replacement[index] != L_ESC {
                    return crate::error!(
                        state.thread,
                        "invalid use of '%c' in replacement string",
                        i32::from(L_ESC)
                    )
                    .map_err(Into::into);
                }

                buffer.push_byte(replacement[index])?;
            } else if replacement[index] == b'0' {
                buffer.push_bytes(&state.src[start..end])?;
            } else {
                push_one_capture(state, (replacement[index] - b'1') as usize, start, end)?;
                buffer.push_stack_value()?;
            }

            index += 1;
        }
    }
    Ok(())
}

/// `add_value`
unsafe fn add_value(
    state: &MatchState<'_, '_>,
    buffer: &mut LuaStringBuilder<'_, '_>,
    start: usize,
    end: usize,
    replacement_type: i32,
) -> VmResult {
    unsafe {
        match replacement_type {
            LUA_TFUNCTION => {
                state.thread.push_value(3)?;
                let captures = push_captures(state, Some((start, end)))?;
                state.thread.call(captures, 1)?;
            }
            LUA_TTABLE => {
                push_one_capture(state, 0, start, end)?;
                let key_slot = state.thread.stack_top().sub(1);
                let key = key_slot.value_unchecked();
                state.thread.get_table_internal(
                    state.thread.to_object(3).unwrap_unchecked(),
                    key,
                    key_slot,
                )?;
            }
            _ => {
                add_s(state, buffer, start, end)?;
                return Ok(());
            }
        }

        if state.thread.to_boolean(-1) == 0 {
            state.thread.pop(1);
            state.thread.push_string(&state.src[start..end])?;
        } else if state.thread.is_string(-1) == 0 {
            let type_name = state.thread.lua_type_name(-1);
            return crate::error!(state.thread, "invalid replacement value (a %s)", &type_name)
                .map_err(Into::into);
        }

        buffer.push_stack_value()?;
    }
    Ok(())
}

/// `str_gsub`
pub(super) fn string_gsub(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let src = thread.check_string(1)?;
        let mut pattern = thread.check_string(2)?;
        let replacement_type = thread.type_of(3);
        let max_s = thread.opt_integer(4, src.len() as i32 + 1)?;
        let anchor = pattern.first() == Some(&b'^');
        let mut substitutions = 0;

        ctx.arg(3).expected(
            matches!(
                replacement_type,
                LUA_TNUMBER | LUA_TSTRING | LUA_TFUNCTION | LUA_TTABLE
            ),
            "string/function/table",
        )?;

        if anchor {
            pattern = &pattern[1..];
        }

        let mut state = new_match_state(thread, src, pattern);
        let mut buffer_storage = LuaStringBuilderStorage::uninit();
        let mut buffer = LuaStringBuilder::new(thread, &mut buffer_storage);
        let mut src_index = 0usize;

        while substitutions < max_s {
            reset_match_state(&mut state);
            let end = match_pattern(&mut state, src_index, 0)?;

            if let Some(end_index) = end {
                substitutions += 1;
                add_value(&state, &mut buffer, src_index, end_index, replacement_type)?;
            }

            if let Some(end_index) = end {
                if end_index > src_index {
                    src_index = end_index;
                } else if src_index < state.src.len() {
                    buffer.push_byte(state.src[src_index])?;
                    src_index += 1;
                } else {
                    break;
                }
            } else if src_index < state.src.len() {
                buffer.push_byte(state.src[src_index])?;
                src_index += 1;
            } else {
                break;
            }

            if anchor {
                break;
            }
        }

        buffer.push_bytes(&state.src[src_index..])?;
        buffer.finish()?;
        thread.push_integer(substitutions)?;
    }
    Ok(2)
}