// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

// Initially generatd by c2rust of 'extern.c' at revision:
// `f14c8ff3f8a164685bc24184fba84904391e378e`.

use std::io;
use std::io::Read;
use std::io::Seek;
use std::io::Write;
use std::mem::MaybeUninit;

use ocamlrep::FromOcamlRep;
use ocamlrep::Header;
use ocamlrep::Value;

use crate::intext::*;

extern "C" {
    fn ocaml_version() -> usize;
}

bitflags::bitflags! {
    /// Flags affecting marshaling
    pub struct ExternFlags: u8 {
        /// Flag to ignore sharing
        const NO_SHARING = 1;
        /// Flag to allow marshaling code pointers. Not permitted in `ocamlrep_marshal`.
        const CLOSURES = 2;
        /// Flag to ensure that output can safely be read back on a 32-bit platform
        const COMPAT_32 = 4;
    }
}

// NB: Must match the definition order in ocaml's marshal.ml
#[derive(ocamlrep::FromOcamlRep)]
enum ExternFlag {
    NoSharing,
    Closures,
    Compat32,
}

impl From<ExternFlag> for ExternFlags {
    fn from(flag: ExternFlag) -> Self {
        match flag {
            ExternFlag::NoSharing => ExternFlags::NO_SHARING,
            ExternFlag::Closures => ExternFlags::CLOSURES,
            ExternFlag::Compat32 => ExternFlags::COMPAT_32,
        }
    }
}

impl FromOcamlRep for ExternFlags {
    fn from_ocamlrep(mut list: ocamlrep::Value<'_>) -> Result<Self, ocamlrep::FromError> {
        let mut res = ExternFlags::empty();
        while !list.is_int() {
            let block = ocamlrep::from::expect_tuple(list, 2)?;
            let flag: ExternFlag = ocamlrep::from::field(block, 0)?;
            res |= flag.into();
            list = block[1];
        }
        Ok(res)
    }
}

// Stack for pending values to marshal

const EXTERN_STACK_INIT_SIZE: usize = 256;
const EXTERN_STACK_MAX_SIZE: usize = 1024 * 1024 * 100;

#[derive(Copy, Clone)]
#[repr(C)]
struct ExternItem<'a> {
    fields: &'a [Value<'a>],
}

// Hash table to record already-marshaled objects and their positions

#[derive(Copy, Clone)]
#[repr(C)]
struct ObjectPosition<'a> {
    obj: Value<'a>,
    pos: usize,
}

// The hash table uses open addressing, linear probing, and a redundant
// representation:
// - a bitvector [present] records which entries of the table are occupied;
// - an array [entries] records (object, position) pairs for the entries
//     that are occupied.
// The bitvector is much smaller than the array (1/128th on 64-bit
// platforms, 1/64th on 32-bit platforms), so it has better locality,
// making it faster to determine that an object is not in the table.
// Also, it makes it faster to empty or initialize a table: only the
// [present] bitvector needs to be filled with zeros, the [entries]
// array can be left uninitialized.

#[repr(C)]
struct PositionTable<'a> {
    shift: u8,
    size: usize,           // size == 1 << (wordsize - shift)
    mask: usize,           // mask == size - 1
    threshold: usize,      // threshold == a fixed fraction of size
    present: Box<[usize]>, // [bitvect_size(size)]
    /// SAFETY: Elements of `entries` are not initialized unless their
    /// corresponding bit is set in `present`.
    entries: Box<[MaybeUninit<ObjectPosition<'a>>]>, // [size]
}

const BITS_WORD: usize = 8 * std::mem::size_of::<usize>();

#[inline]
const fn bitvect_size(n: usize) -> usize {
    (n + BITS_WORD - 1) / BITS_WORD
}

const POS_TABLE_INIT_SIZE_LOG2: usize = 8;
const POS_TABLE_INIT_SIZE: usize = 1 << POS_TABLE_INIT_SIZE_LOG2;

// Multiplicative Fibonacci hashing
// (Knuth, TAOCP vol 3, section 6.4, page 518).
// HASH_FACTOR is (sqrt(5) - 1) / 2 * 2^wordsize.
const HASH_FACTOR: usize = 11400714819323198486;
#[inline]
const fn hash(v: Value<'_>, shift: u8) -> usize {
    v.to_bits().wrapping_mul(HASH_FACTOR) >> shift
}

// When the table becomes 2/3 full, its size is increased.
#[inline]
const fn threshold(sz: usize) -> usize {
    (sz * 2) / 3
}

// Accessing bitvectors

#[inline]
fn bitvect_test(bv: &[usize], i: usize) -> bool {
    bv[i / BITS_WORD] & (1 << (i & (BITS_WORD - 1))) != 0
}

#[inline]
fn bitvect_set(bv: &mut [usize], i: usize) {
    bv[i / BITS_WORD] |= 1 << (i & (BITS_WORD - 1));
}

// Conversion to big-endian

#[inline]
fn store16(dst: &mut impl Write, n: i16) -> io::Result<()> {
    dst.write_all(&n.to_be_bytes())
}

#[inline]
fn store32(dst: &mut impl Write, n: i32) -> io::Result<()> {
    dst.write_all(&n.to_be_bytes())
}

#[inline]
fn store64(dst: &mut impl Write, n: i64) -> io::Result<()> {
    dst.write_all(&n.to_be_bytes())
}

#[repr(C)]
struct State<'a, W: Write> {
    flags: ExternFlags, // logical or of some of the flags

    obj_counter: usize, // Number of objects emitted so far
    size_32: usize,     // Size in words of 32-bit block for struct.
    size_64: usize,     // Size in words of 64-bit block for struct.

    // Stack for pending value to marshal
    stack: Vec<ExternItem<'a>>,

    // Hash table to record already marshalled objects
    pos_table: PositionTable<'a>,

    // The output file or buffer we are writing to
    output: W,
    output_len: usize,
}

impl<'a, W: Write> State<'a, W> {
    fn new(output: W) -> Self {
        Self {
            flags: ExternFlags::empty(),
            obj_counter: 0,
            size_32: 0,
            size_64: 0,

            stack: Vec::with_capacity(EXTERN_STACK_INIT_SIZE),

            pos_table: PositionTable {
                shift: 0,
                size: 0,
                mask: 0,
                threshold: 0,
                present: [].into(),
                entries: [].into(),
            },

            output,
            output_len: 0,
        }
    }

    /// Initialize the position table
    fn init_position_table(&mut self) {
        if self.flags.contains(ExternFlags::NO_SHARING) {
            return;
        }
        self.pos_table.size = POS_TABLE_INIT_SIZE;
        self.pos_table.shift =
            (8 * std::mem::size_of::<Value<'a>>() - POS_TABLE_INIT_SIZE_LOG2) as u8;
        self.pos_table.mask = POS_TABLE_INIT_SIZE - 1;
        self.pos_table.threshold = threshold(POS_TABLE_INIT_SIZE);
        // SAFETY: zero is a valid value for the elements of `present`.
        unsafe {
            self.pos_table.present =
                Box::new_zeroed_slice(bitvect_size(POS_TABLE_INIT_SIZE)).assume_init();
        }
        self.pos_table.entries = Box::new_uninit_slice(POS_TABLE_INIT_SIZE);
    }

    /// Grow the position table
    fn resize_position_table(&mut self) {
        let new_size: usize;
        let new_shift: u8;

        // Grow the table quickly (x 8) up to 10^6 entries,
        // more slowly (x 2) afterwards.
        if self.pos_table.size < 1000000 {
            new_size = self.pos_table.size * 8;
            new_shift = self.pos_table.shift - 3;
        } else {
            new_size = self.pos_table.size * 2;
            new_shift = self.pos_table.shift - 1;
        }
        let old = std::mem::replace(
            &mut self.pos_table,
            PositionTable {
                size: new_size,
                shift: new_shift,
                mask: new_size - 1,
                threshold: threshold(new_size),
                // SAFETY: zero is a valid value for the elements of `present`.
                present: unsafe { Box::new_zeroed_slice(bitvect_size(new_size)).assume_init() },
                entries: Box::new_uninit_slice(new_size),
            },
        );

        // Insert every entry of the old table in the new table
        let mut i = 0;
        while i < old.size {
            if bitvect_test(&old.present, i) {
                // SAFETY: We checked that the bit for `i` is set in
                // `old.present`, so `entries[i]` must be initialized
                let old_entry = unsafe { old.entries[i].assume_init() };
                let mut h = hash(old_entry.obj, self.pos_table.shift);
                while bitvect_test(&self.pos_table.present, h) {
                    h = (h + 1) & self.pos_table.mask
                }
                bitvect_set(&mut self.pos_table.present, h);
                self.pos_table.entries[h] = MaybeUninit::new(old_entry);
            }
            i += 1
        }
    }

    /// Determine whether the given object [obj] is in the hash table.
    /// If so, set `*pos_out` to its position in the output and return true.
    /// If not, set `*h_out` to the hash value appropriate for
    /// `record_location` and return false.
    #[inline]
    fn lookup_position(&self, obj: Value<'a>, pos_out: &mut usize, h_out: &mut usize) -> bool {
        let mut h: usize = hash(obj, self.pos_table.shift);
        loop {
            if !bitvect_test(&self.pos_table.present, h) {
                *h_out = h;
                return false;
            }
            // SAFETY: We checked that the bit for `h` is set in `present`, so
            // `entries[h]` must be initialized
            let entry = unsafe { self.pos_table.entries[h].assume_init_ref() };
            if entry.obj == obj {
                *pos_out = entry.pos;
                return true;
            }
            h = (h + 1) & self.pos_table.mask
        }
    }

    /// Record the output position for the given object [obj].
    ///
    /// The [h] parameter is the index in the hash table where the object
    /// must be inserted.  It was determined during lookup.
    fn record_location(&mut self, obj: Value<'a>, h: usize) {
        if self.flags.contains(ExternFlags::NO_SHARING) {
            return;
        }
        bitvect_set(&mut self.pos_table.present, h);
        self.pos_table.entries[h] = MaybeUninit::new(ObjectPosition {
            obj,
            pos: self.obj_counter,
        });
        self.obj_counter += 1;
        if self.obj_counter >= self.pos_table.threshold {
            self.resize_position_table();
        };
    }

    // Write characters, integers, and blocks in the output buffer

    #[inline]
    fn write(&mut self, c: u8) -> io::Result<()> {
        self.output_len += 1;
        self.output.write_all(&[c])
    }

    fn writeblock(&mut self, data: &[u8]) -> io::Result<()> {
        self.output_len += data.len();
        self.output.write_all(data)
    }

    #[inline]
    fn writeblock_float8(&mut self, data: &[f64]) -> io::Result<()> {
        if ARCH_FLOAT_ENDIANNESS == 0x01234567 || ARCH_FLOAT_ENDIANNESS == 0x76543210 {
            // SAFETY: `data.as_ptr()` will be valid for reads of `data.len() *
            // size_of::<f64>()` bytes
            self.writeblock(unsafe {
                std::slice::from_raw_parts(
                    data.as_ptr() as *const u8,
                    data.len() * std::mem::size_of::<f64>(),
                )
            })
        } else {
            unimplemented!()
        }
    }

    fn writecode8(&mut self, code: u8, val: i8) -> io::Result<()> {
        self.output_len += 2;
        self.output.write_all(&[code, val as u8])
    }

    fn writecode16(&mut self, code: u8, val: i16) -> io::Result<()> {
        self.output_len += 3;
        self.output.write_all(&[code])?;
        store16(&mut self.output, val)
    }

    fn writecode32(&mut self, code: u8, val: i32) -> io::Result<()> {
        self.output_len += 5;
        self.output.write_all(&[code])?;
        store32(&mut self.output, val)
    }

    fn writecode64(&mut self, code: u8, val: i64) -> io::Result<()> {
        self.output_len += 9;
        self.output.write_all(&[code])?;
        store64(&mut self.output, val)
    }

    /// Marshaling integers
    #[inline]
    fn extern_int(&mut self, n: isize) -> io::Result<()> {
        if (0..0x40).contains(&n) {
            self.write(PREFIX_SMALL_INT + n as u8)
        } else if (-(1 << 7)..(1 << 7)).contains(&n) {
            self.writecode8(CODE_INT8, n as i8)
        } else if (-(1 << 15)..(1 << 15)).contains(&n) {
            self.writecode16(CODE_INT16, n as i16)
        } else if !(-(1 << 30)..(1 << 30)).contains(&n) {
            if self.flags.contains(ExternFlags::COMPAT_32) {
                panic!("output_value: integer cannot be read back on 32-bit platform");
            }
            self.writecode64(CODE_INT64, n as i64)
        } else {
            self.writecode32(CODE_INT32, n as i32)
        }
    }

    /// Marshaling references to previously-marshaled blocks
    #[inline]
    fn extern_shared_reference(&mut self, d: usize) -> io::Result<()> {
        if d < 0x100 {
            self.writecode8(CODE_SHARED8, d as i8)
        } else if d < 0x10000 {
            self.writecode16(CODE_SHARED16, d as i16)
        } else if d >= 1 << 32 {
            self.writecode64(CODE_SHARED64, d as i64)
        } else {
            self.writecode32(CODE_SHARED32, d as i32)
        }
    }

    /// Marshaling block headers
    #[inline]
    fn extern_header(&mut self, sz: usize, tag: u8) -> io::Result<()> {
        if tag < 16 && sz < 8 {
            self.write(PREFIX_SMALL_BLOCK + tag + ((sz as u8) << 4))
        } else {
            // Note: ocaml-14.4.0 uses `Caml_white` (`0 << 8`)
            // ('caml/runtime/gc.h').
            //
            // In ocaml-5, via PR https://github.com/ocaml/ocaml/pull/10831, in
            // commit
            // `https://github.com/ocaml/ocaml/commit/868265f4532a2cc33bbffd83221c9613e743d759`
            // this becomes,
            //   let hd: header_t = Make_header(sz, tag, NOT_MARKABLE);
            // where, `NOT_MARKABLE` (`3 << 8`) ('caml/runtime/shared_heap.h').

            // Check the prevailing OCaml version is well initialized & one
            // we've tested for.
            let which_ocaml = unsafe { ocaml_version() };
            if ![41400, 50000].contains(&which_ocaml) {
                panic!("unexpected ocaml version: {which_ocaml}!");
            }

            let color = if which_ocaml < 50000 {
                ocamlrep::Color::White
            } else {
                ocamlrep::Color::Black
            };
            let hd = Header::with_color(sz, tag, color).to_bits();

            if sz > 0x3FFFFF && self.flags.contains(ExternFlags::COMPAT_32) {
                panic!("output_value: array cannot be read back on 32-bit platform");
            }
            if hd < 1 << 32 {
                self.writecode32(CODE_BLOCK32, hd as i32)
            } else {
                self.writecode64(CODE_BLOCK64, hd as i64)
            }
        }
    }

    #[inline]
    fn extern_string(&mut self, bytes: &'a [u8]) -> io::Result<()> {
        let len = bytes.len();
        if len < 0x20 {
            self.write(PREFIX_SMALL_STRING + len as u8)?;
        } else if len < 0x100 {
            self.writecode8(CODE_STRING8, len as i8)?;
        } else {
            if len > 0xFFFFFB && self.flags.contains(ExternFlags::COMPAT_32) {
                panic!("output_value: string cannot be read back on 32-bit platform");
            }
            if len < 1 << 32 {
                self.writecode32(CODE_STRING32, len as i32)?;
            } else {
                self.writecode64(CODE_STRING64, len as i64)?;
            }
        }
        self.writeblock(bytes)
    }

    /// Marshaling FP numbers
    #[inline]
    fn extern_double(&mut self, v: f64) -> io::Result<()> {
        self.write(CODE_DOUBLE_NATIVE)?;
        self.writeblock_float8(&[v])
    }

    /// Marshaling FP arrays
    #[inline]
    fn extern_double_array(&mut self, slice: &[f64]) -> io::Result<()> {
        let nfloats = slice.len();
        if nfloats < 0x100 {
            self.writecode8(CODE_DOUBLE_ARRAY8_NATIVE, nfloats as i8)?;
        } else {
            if nfloats > 0x1FFFFF && self.flags.contains(ExternFlags::COMPAT_32) {
                panic!("output_value: float array cannot be read back on 32-bit platform");
            }
            if nfloats < 1 << 32 {
                self.writecode32(CODE_DOUBLE_ARRAY32_NATIVE, nfloats as i32)?;
            } else {
                self.writecode64(CODE_DOUBLE_ARRAY64_NATIVE, nfloats as i64)?;
            }
        }
        self.writeblock_float8(slice)
    }

    /// Marshal the given value in the output buffer
    fn extern_rec(&mut self, mut v: Value<'a>) -> io::Result<()> {
        let mut goto_next_item: bool;

        let mut h: usize = 0;
        let mut pos: usize = 0;

        self.init_position_table();

        loop {
            if v.is_int() {
                self.extern_int(v.as_int().unwrap())?;
            } else {
                let b = v.as_block().unwrap();
                let tag = b.tag();
                let sz = b.size();

                if tag == ocamlrep::FORWARD_TAG {
                    let f = b[0];
                    if let Some(fb) = f.as_block() {
                        if fb.tag() == ocamlrep::FORWARD_TAG
                            || fb.tag() == ocamlrep::LAZY_TAG
                            || fb.tag() == ocamlrep::FORCING_TAG
                            || fb.tag() == ocamlrep::DOUBLE_TAG
                        {
                            // Do not short-circuit the pointer.
                        } else {
                            v = f;
                            continue;
                        }
                    } else {
                        v = f;
                        continue;
                    }
                }
                // Atoms are treated specially for two reasons: they are not allocated
                // in the externed block, and they are automatically shared.
                if sz == 0 {
                    self.extern_header(0, tag)?;
                } else {
                    // Check if object already seen
                    if !self.flags.contains(ExternFlags::NO_SHARING) {
                        if self.lookup_position(v, &mut pos, &mut h) {
                            self.extern_shared_reference(self.obj_counter - pos)?;
                            goto_next_item = true;
                        } else {
                            goto_next_item = false;
                        }
                    } else {
                        goto_next_item = false;
                    }
                    if !goto_next_item {
                        // Output the contents of the object
                        match tag {
                            ocamlrep::STRING_TAG => {
                                let bytes = v.as_byte_string().unwrap();
                                let len: usize = bytes.len();
                                self.extern_string(bytes)?;
                                self.size_32 += 1 + (len + 4) / 4;
                                self.size_64 += 1 + (len + 8) / 8;
                                self.record_location(v, h);
                            }
                            ocamlrep::DOUBLE_TAG => {
                                self.extern_double(v.as_float().unwrap())?;
                                self.size_32 += 1 + 2;
                                self.size_64 += 1 + 1;
                                self.record_location(v, h);
                            }
                            ocamlrep::DOUBLE_ARRAY_TAG => {
                                let slice = v.as_double_array().unwrap();
                                self.extern_double_array(slice)?;
                                let nfloats = slice.len();
                                self.size_32 += 1 + nfloats * 2;
                                self.size_64 += 1 + nfloats;
                                self.record_location(v, h);
                            }
                            ocamlrep::ABSTRACT_TAG => {
                                panic!("output_value: abstract value (Abstract)");
                            }
                            // INFIX_TAG represents an infix header inside a
                            // closure, and can only occur in blocks with tag
                            // CLOSURE_TAG
                            ocamlrep::INFIX_TAG => {
                                panic!("output_value: marshaling of closures not implemented");
                            }
                            ocamlrep::CUSTOM_TAG => {
                                panic!("output_value: marshaling of custom blocks not implemented");
                            }
                            ocamlrep::CLOSURE_TAG => {
                                panic!("output_value: marshaling of closures not implemented");
                            }
                            _ => {
                                self.extern_header(sz, tag)?;
                                self.size_32 += 1 + sz;
                                self.size_64 += 1 + sz;
                                self.record_location(v, h);
                                // Remember that we still have to serialize fields 1 ... sz - 1
                                if sz > 1 {
                                    if self.stack.len() + 1 >= EXTERN_STACK_MAX_SIZE {
                                        panic!("Stack overflow in marshaling value");
                                    }
                                    self.stack.push(ExternItem {
                                        fields: &b.as_values().unwrap()[1..],
                                    });
                                }
                                // Continue serialization with the first field
                                v = v.field(0).unwrap();
                                continue;
                            }
                        }
                    }
                }
            }
            // C goto label `next_item:` here

            // Pop one more item to marshal, if any
            if let Some(item) = self.stack.last_mut() {
                v = item.fields[0];
                item.fields = &item.fields[1..];
                if item.fields.is_empty() {
                    self.stack.pop();
                }
            } else {
                // We are done.
                return Ok(());
            }
        }
        // Never reached as function leaves with return
    }

    fn extern_value(
        &mut self,
        v: Value<'a>,
        flags: ExternFlags,
        mut header: &mut [u8],  // out
        header_len: &mut usize, // out
    ) -> io::Result<usize> {
        // Initializations
        self.flags = flags;
        self.obj_counter = 0;
        self.size_32 = 0;
        self.size_64 = 0;
        // Marshal the object
        self.extern_rec(v)?;
        // Write the header
        let res_len = self.output_len;
        if res_len >= (1 << 32) || self.size_32 >= (1 << 32) || self.size_64 >= (1 << 32) {
            // The object is too big for the small header format.
            // Fail if we are in compat32 mode, or use big header.
            if self.flags.contains(ExternFlags::COMPAT_32) {
                panic!("output_value: object too big to be read back on 32-bit platform");
            }
            store32(&mut header, MAGIC_NUMBER_BIG as i32)?;
            store32(&mut header, 0)?;
            store64(&mut header, res_len as i64)?;
            store64(&mut header, self.obj_counter as i64)?;
            store64(&mut header, self.size_64 as i64)?;
            *header_len = 32;
            Ok(res_len)
        } else {
            // Use the small header format
            store32(&mut header, MAGIC_NUMBER_SMALL as i32)?;
            store32(&mut header, res_len as i32)?;
            store32(&mut header, self.obj_counter as i32)?;
            store32(&mut header, self.size_32 as i32)?;
            store32(&mut header, self.size_64 as i32)?;
            *header_len = 20;
            Ok(res_len)
        }
    }
}

pub fn output_value<W: Read + Write + Seek>(
    w: &mut W,
    v: Value<'_>,
    flags: ExternFlags,
) -> io::Result<()> {
    let mut header = [0; 32];
    let mut header_len = 0;
    // At this point we don't know the size of the header.
    // Guess that it is small, and fix up later if not.
    w.rewind()?;
    w.write_all(&[0; 20])?;
    let mut s = State::new(&mut *w);
    s.extern_value(v, flags, &mut header, &mut header_len)?;
    drop(s);
    w.flush()?;
    if header_len != 20 {
        // Bad guess! Need to shift the output to make room for big header.
        w.seek(std::io::SeekFrom::Start(20))?;
        let mut output = vec![];
        w.read_to_end(&mut output)?;
        w.seek(std::io::SeekFrom::Start(header_len as u64))?;
        w.write_all(&output)?;
    }
    w.rewind()?;
    w.write_all(&header[0..header_len])?;
    w.flush()
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.
#![feature(exit_status_error)]

use ocamlrep::FromOcamlRep;

type OcamlValue = usize;

#[no_mangle]
unsafe extern "C" fn ocamlrep_marshal_output_value_to_string(
    v: OcamlValue,
    flags: OcamlValue,
) -> OcamlValue {
    ocamlrep_ocamlpool::catch_unwind(|| {
        let v = ocamlrep::Value::from_bits(v);
        let flags = ocamlrep_marshal::ExternFlags::from_ocaml(flags).unwrap();
        let mut cursor = std::io::Cursor::new(vec![]);
        ocamlrep_marshal::output_value(&mut cursor, v, flags).unwrap();
        ocamlrep_ocamlpool::to_ocaml(&cursor.into_inner())
    })
}

#[no_mangle]
unsafe extern "C" fn ocamlrep_marshal_input_value_from_string(
    str: OcamlValue,
    ofs: OcamlValue,
) -> OcamlValue {
    ocamlrep_ocamlpool::catch_unwind(|| {
        let offset = usize::from_ocaml(ofs).unwrap();
        let str = ocamlrep::bytes_from_ocamlrep(ocamlrep::Value::from_bits(str)).unwrap();
        let str = &str[offset..];
        let pool = ocamlrep_ocamlpool::Pool::new();
        ocamlrep_marshal::input_value(str, &pool).to_bits()
    })
}

// Hack! Trick buck into believing that these libraries are used. See [Note:
// Test blocks for Cargo] in `ocamlrep_ocamlpool/test/ocamlpool_test.rs`.
const _: () = {
    #[allow(unused_imports)]
    use anyhow;
    #[allow(unused_imports)]
    use cargo_test_utils;
    #[allow(unused_imports)]
    use tempdir;
};

#[cfg(test)]
mod tests {
    use anyhow::Result;
    use cargo_test_utils::*;
    use tempdir::TempDir;

    #[test]
    fn ocamlrep_marshal_test() -> Result<()> {
        let tmp_dir = TempDir::new("ocamlrep_marshal_test")?;
        std::fs::copy(
            "test_ocamlrep_marshal.ml",
            tmp_dir.path().join("test_ocamlrep_marshal.ml"),
        )?;
        let compile_cmd = cmd(
            "ocamlopt.opt",
            &[
                "-verbose",
                "-c",
                "test_ocamlrep_marshal.ml",
                "-o",
                "test_ocamlrep_marshal_ml.cmx",
            ],
            Some(tmp_dir.path()),
        );
        assert_eq!(run(compile_cmd).map_err(fmt_exit_status_err), Ok(()));
        let link_cmd = cmd(
            "ocamlopt.opt",
            &[
                "-verbose",
                "-o",
                "ocamlrep_marshal_test",
                "test_ocamlrep_marshal_ml.cmx",
                "-ccopt",
                &("-L".to_owned() + workspace_dir(&["target", build_flavor()]).to_str().unwrap()),
                "-cclib",
                "-locamlrep_marshal",
                "-cclib",
                "-locamlrep_marshal_ffi_bindings",
                "-cclib",
                "-locamlrep_ocamlpool",
            ],
            Some(tmp_dir.path()),
        );
        assert_eq!(run(link_cmd).map_err(fmt_exit_status_err), Ok(()));
        let ocamlrep_marshal_test_cmd = cmd(
            tmp_dir
                .path()
                .join("ocamlrep_marshal_test")
                .as_path()
                .to_str()
                .unwrap(),
            &[],
            None,
        );
        assert_eq!(
            run(ocamlrep_marshal_test_cmd).map_err(fmt_exit_status_err),
            Ok(())
        );
        tmp_dir.close()?;
        Ok(())
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.

// Assume an opam environment (`eval "$(opam env --switch=default
// --set-switch)"`) then to find the prevailing standard library caml
// headers, `OCAMLLIB=$(ocamlopt.opt -config | grep standard_library:
// | awk '{ print $2 }')`.
fn ocamllib_dir() -> std::path::PathBuf {
    let mut sh = std::process::Command::new("sh");
    sh.args([
        "-c",
        "ocamlopt.opt -config | grep standard_library: | awk '{ print $2 }'",
    ]);
    std::path::Path::new(
        std::str::from_utf8(&sh.output().unwrap().stdout)
            .unwrap()
            .trim(),
    )
    .to_path_buf()
}

fn main() {
    // Tell Cargo that if the given file changes, to rerun this build script.
    println!("cargo:rerun-if-changed=ocamlpool.c");
    cc::Build::new()
        .include(ocamllib_dir().as_path().to_str().unwrap())
        .file("ocamlpool.c")
        .compile("ocamlpool");
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use std::ffi::CString;
use std::panic::UnwindSafe;

pub use bumpalo::Bump;
use ocamlrep::Allocator;
use ocamlrep::BlockBuilder;
pub use ocamlrep::FromOcamlRep;
pub use ocamlrep::FromOcamlRepIn;
use ocamlrep::MemoizationCache;
use ocamlrep::ToOcamlRep;
pub use ocamlrep::Value;

extern "C" {
    fn ocamlpool_enter();
    fn ocamlpool_leave();
    fn ocamlpool_reserve_block(tag: u8, size: usize) -> usize;
    fn caml_failwith(msg: *const i8);
    fn caml_initialize(addr: *mut usize, value: usize);
    static ocamlpool_generation: usize;

    pub fn caml_named_value(name: *const std::ffi::c_char) -> *mut usize;
    pub fn caml_callbackN_exn(closure: usize, n: std::ffi::c_int, args: *const usize) -> usize;
}

pub struct Pool {
    cache: MemoizationCache,
}

impl Pool {
    /// Prepare the ocamlpool library to allocate values directly on the OCaml
    /// runtime's garbage-collected heap.
    ///
    /// # Safety
    ///
    /// The OCaml runtime is not thread-safe, and this function will interact
    /// with it. If any other thread interacts with the OCaml runtime or
    /// ocamlpool library during the lifetime of the `Pool`, undefined behavior
    /// will result.
    #[inline(always)]
    pub unsafe fn new() -> Self {
        ocamlpool_enter();
        Self {
            cache: MemoizationCache::new(),
        }
    }

    #[inline(always)]
    pub fn add<'a, T: ToOcamlRep + ?Sized>(&'a self, value: &'a T) -> Value<'a> {
        value.to_ocamlrep(self)
    }
}

impl Drop for Pool {
    #[inline(always)]
    fn drop(&mut self) {
        unsafe { ocamlpool_leave() };
    }
}

impl Allocator for Pool {
    #[inline(always)]
    fn generation(&self) -> usize {
        unsafe { ocamlpool_generation }
    }

    #[inline(always)]
    fn block_with_size_and_tag(&self, size: usize, tag: u8) -> BlockBuilder<'_> {
        let ptr = unsafe { ocamlpool_reserve_block(tag, size) as *mut Value<'_> };
        BlockBuilder::new(unsafe { std::slice::from_raw_parts_mut(ptr, size) })
    }

    #[inline(always)]
    fn set_field<'a>(&self, block: &mut BlockBuilder<'a>, index: usize, value: Value<'a>) {
        assert!(index < block.size());
        unsafe {
            caml_initialize(
                (self.block_ptr_mut(block) as *mut usize).add(index),
                value.to_bits(),
            )
        };
    }

    unsafe fn block_ptr_mut<'a>(&self, block: &mut BlockBuilder<'a>) -> *mut Value<'a> {
        block.address() as *mut _
    }

    fn memoized<'a>(
        &'a self,
        ptr: usize,
        size: usize,
        f: impl FnOnce(&'a Self) -> Value<'a>,
    ) -> Value<'a> {
        let bits = self.cache.memoized(ptr, size, || f(self).to_bits());
        // SAFETY: The only memoized values in the cache are those computed in
        // the closure on the previous line. Since f returns Value<'a>, any
        // cached bits must represent a valid Value<'a>,
        unsafe { Value::from_bits(bits) }
    }

    fn add_root<'a, T: ToOcamlRep + ?Sized>(&'a self, value: &'a T) -> Value<'a> {
        self.cache.with_cache(|| value.to_ocamlrep(self))
    }
}

/// Convert the given value to an OCaml value on the OCaml runtime's
/// garbage-collected heap.
///
/// # Safety
///
/// The OCaml runtime is not thread-safe, and this function will interact with
/// it. If any other thread interacts with the OCaml runtime or ocamlpool
/// library during the execution of `to_ocaml`, undefined behavior will result.
///
/// # Panics
///
/// Panics upon attempts to re-enter `to_ocaml`.
#[inline(always)]
pub unsafe fn to_ocaml<T: ToOcamlRep + ?Sized>(value: &T) -> usize {
    let pool = Pool::new();
    let result = pool.add_root(value);
    result.to_bits()
}

/// Catches panics in `f` and raises a OCaml exception of type Failure
/// with the panic message (if the panic was raised with a `&str` or `String`).
pub fn catch_unwind(f: impl FnOnce() -> usize + UnwindSafe) -> usize {
    catch_unwind_with_handler(f, |msg: &str| -> Result<usize, String> { Err(msg.into()) })
}

/// Catches panics in `f` and raises a OCaml exception of type Failure
/// with the panic message (if the panic was raised with a `&str` or `String`).
/// `h` handles panic msg, it may re-raise by returning Err.
pub fn catch_unwind_with_handler(
    f: impl FnOnce() -> usize + UnwindSafe,
    h: impl FnOnce(&str) -> Result<usize, String>,
) -> usize {
    let err = match std::panic::catch_unwind(f) {
        Ok(value) => return value,
        Err(err) => err,
    };
    let msg: &str = if let Some(s) = err.downcast_ref::<&str>() {
        s
    } else if let Some(s) = err.downcast_ref::<String>() {
        s.as_str()
    } else {
        // TODO: Build a smarter message in this case (using panic::set_hook?)
        "Panicked with non-string object"
    };
    match h(msg) {
        Ok(value) => return value,
        Err(err) => unsafe {
            let msg = CString::new(err).unwrap();
            caml_failwith(msg.as_ptr().cast());
        },
    }
    unreachable!();
}

/// Assume that some Pool exists in some parent scope. Since ocamlpool is
/// implemented with statics, we don't need a reference to that pool to write to
/// it.
///
/// Does not preserve sharing of values referred to by multiple references or
/// Rcs (but sharing is preserved for `ocamlrep::rc::RcOc`).
///
/// # Safety
///
/// The OCaml runtime is not thread-safe, and this function will interact with
/// it. If any other thread interacts with the OCaml runtime or ocamlpool
/// library during the execution of this function, undefined behavior will
/// result.
#[inline(always)]
pub unsafe fn add_to_ambient_pool<T: ToOcamlRep>(value: &T) -> usize {
    let fake_pool = Pool {
        cache: MemoizationCache::new(),
    };
    let result = value.to_ocamlrep(&fake_pool).to_bits();
    std::mem::forget(fake_pool);
    result
}

/// Check if an OCaml value is an exception.
///
/// For internal purposes.
pub fn is_exception_result(v: usize) -> bool {
    v & 3 == 2
}

#[macro_export]
macro_rules! ocaml_ffi_fn {
    (fn $name:ident($($param:ident: $ty:ty),+  $(,)?) -> $ret:ty $code:block) => {
        #[no_mangle]
        pub unsafe extern "C" fn $name ($($param: usize,)*) -> usize {
            $crate::catch_unwind(|| {
                fn inner($($param: $ty,)*) -> $ret { $code }
                use $crate::FromOcamlRep;
                $(let $param = <$ty>::from_ocaml($param).unwrap();)*
                let result = inner($($param,)*);
                $crate::to_ocaml(&result)
            })
        }
    };

    (fn $name:ident() -> $ret:ty $code:block) => {
        #[no_mangle]
        pub unsafe extern "C" fn $name (_unit: usize) -> usize {
            $crate::catch_unwind(|| {
                fn inner() -> $ret { $code }
                let result = inner();
                $crate::to_ocaml(&result)
            })
        }
    };

    (fn $name:ident($($param:ident: $ty:ty),*  $(,)?) $code:block) => {
        $crate::ocaml_ffi_fn! {
            fn $name($($param: $ty),*) -> () $code
        }
    };
}

/// Convenience macro for declaring OCaml FFI wrappers.
///
/// Each parameter will be converted from OCaml using `ocamlrep` and the result
/// will be converted to OCaml using `ocamlrep` and allocated on the OCaml GC
/// heap using `ocamlpool`.
///
/// Panics in the function body will be caught and converted to an OCaml
/// exception of type Failure.
#[macro_export]
macro_rules! ocaml_ffi {
    ($(fn $name:ident($($param:ident: $ty:ty),*  $(,)?) $(-> $ret:ty)? $code:block)*) => {
        $($crate::ocaml_ffi_fn! {
            fn $name($($param: $ty),*) $(-> $ret)* $code
        })*
    };
}

#[macro_export]
macro_rules! ocaml_ffi_with_arena_fn {
    (fn $name:ident<$lifetime:lifetime>($arena:ident: $arena_ty:ty, $($param:ident: $ty:ty),+ $(,)?) -> $ret:ty $code:block) => {
        #[no_mangle]
        pub unsafe extern "C" fn $name ($($param: usize,)*) -> usize {
            $crate::catch_unwind(|| {
                use $crate::FromOcamlRepIn;
                let arena = &$crate::Bump::new();
                fn inner<$lifetime>($arena: $arena_ty, $($param: usize,)*) -> $ret {
                    $(let $param = unsafe {
                        <$ty>::from_ocamlrep_in($crate::Value::from_bits($param), $arena).unwrap()
                    };)*
                    $code
                }
                let result = inner(arena, $($param,)*);
                $crate::to_ocaml(&result)
            })
        }
    };

    (fn $name:ident<$lifetime:lifetime>($arena:ident: $arena_ty:ty $(,)?) -> $ret:ty $code:block) => {
        #[no_mangle]
        pub unsafe extern "C" fn $name (_unit: usize) -> usize {
            $crate::catch_unwind(|| {
                fn inner<$lifetime>($arena: $arena_ty) -> $ret { $code }
                let arena = &$crate::Bump::new();
                let result = inner(arena);
                $crate::to_ocaml(&result)
            })
        }
    };

    (fn $name:ident<$lifetime:lifetime>($($param:ident: $ty:ty),* $(,)?) $code:block) => {
        $crate::ocaml_ffi_with_arena_fn! {
            fn $name<$lifetime>($($param: $ty),*) -> () $code
        }
    };
}

/// Convenience macro for declaring OCaml FFI wrappers which use an arena to
/// allocate the arguments and return value.
///
/// FFI functions declared with this macro must declare exactly one lifetime
/// parameter. The function's first value parameter must be a reference to a
/// `bumpalo::Bump` arena with that lifetime:
///
/// ```
/// ocaml_ffi_with_arena! {
///     fn swap_str_pair<'a>(arena: &'a Bump, pair: (&'a str, &'a str)) -> (&'a str, &'a str) {
///         (pair.1, pair.0)
///     }
/// }
/// ```
///
/// An OCaml extern declaration for this function would look like this:
///
/// ```
/// external swap_str_pair : string * string -> string * string = "swap_str_pair"
/// ```
///
/// Note that no parameter for the arena appears on the OCaml side--it is
/// constructed on the Rust side and lives only for the duration of one FFI
/// call.
///
/// Each (non-arena) parameter will be converted from OCaml using
/// `ocamlrep::FromOcamlRepIn`, and allocated in the given arena (if its
/// `FromOcamlRepIn` implementation makes use of the arena).
///
/// The return value (which may be allocated in the given arena, if convenient)
/// will be converted to OCaml using `ocamlrep::ToOcamlRep`. The converted OCaml
/// value will be allocated on the OCaml heap using `ocamlpool`.
///
/// Panics in the function body will be caught and converted to an OCaml
/// exception of type `Failure`.
#[macro_export]
macro_rules! ocaml_ffi_with_arena {
    ($(fn $name:ident<$lifetime:lifetime>($($param:ident: $ty:ty),* $(,)?) $(-> $ret:ty)? $code:block)*) => {
        $($crate::ocaml_ffi_with_arena_fn! {
            fn $name<$lifetime>($($param: $ty),*) $(-> $ret)* $code
        })*
    };
}

#[macro_export]
macro_rules! ocaml_ffi_arena_result_fn {
    (fn $name:ident<$lifetime:lifetime>($arena:ident: $arena_ty:ty, $($param:ident: $ty:ty),+ $(,)?) -> $ret:ty $code:block) => {
        #[no_mangle]
        pub unsafe extern "C" fn $name ($($param: usize,)*) -> usize {
            $crate::catch_unwind(|| {
                fn inner<$lifetime>($arena: $arena_ty, $($param: $ty,)*) -> $ret {
                    $code
                }
                use $crate::FromOcamlRep;
                $(let $param = <$ty>::from_ocaml($param).unwrap();)*
                let arena = &$crate::Bump::new();
                let result = inner(arena, $($param,)*);
                $crate::to_ocaml(&result)
            })
        }
    };
}

/// Convenience macro for declaring OCaml FFI wrappers which use an arena to
/// just the return value.
///
/// FFI functions declared with this macro must declare exactly one lifetime
/// parameter. The function's first value parameter must be a reference to a
/// `bumpalo::Bump` arena with that lifetime:
///
/// ```
/// ocaml_ffi_arena_result! {
///     fn parse_example<'a>(arena: &'a Bump, text: String) -> ('a str, &'a str) {
///         /* copy parts of text into arena while parsing */
///         (pair.1, pair.0)
///     }
/// }
/// ```
///
/// An OCaml extern declaration for this function would look like this:
///
/// ```
/// external parse_example : string -> string * string = "parse_example"
/// ```
///
/// Note that no parameter for the arena appears on the OCaml side--it is
/// constructed on the Rust side and lives only for the duration of one FFI
/// call.
///
/// Each parameter after the arena parameter will be converted from OCaml using
/// `ocamlrep` and passed in as an owned value.
///
/// The return value (which may be allocated in the given arena, if convenient)
/// will be converted to OCaml using `ocamlrep::ToOcamlRep`. The converted OCaml
/// value will be allocated on the OCaml heap using `ocamlpool`.
///
/// Panics in the function body will be caught and converted to an OCaml
/// exception of type `Failure`.
#[macro_export]
macro_rules! ocaml_ffi_arena_result {
    ($(fn $name:ident<$lifetime:lifetime>($($param:ident: $ty:ty),* $(,)?) $(-> $ret:ty)? $code:block)*) => {
        $($crate::ocaml_ffi_arena_result_fn! {
            fn $name<$lifetime>($($param: $ty),*) $(-> $ret)* $code
        })*
    };
}

#[macro_export]
macro_rules! ocaml_registered_function_fn {
    // This needs to be first, as macro matching is linear.
    //
    // caml_callback_exn works as it directly calls into the implemented OCaml functions.
    //
    // caml_callback{2,3,N}_exn don't work, as they go through caml_apply2,
    // caml_apply3 etc. which for some reason crashes!
    //
    // TODO: FIgure out how to make caml_apply2 and friends not crash, and remove the below rule.
    ($ocaml_name:expr, fn $name:ident($param1:ident: $ty1:ty, $($params:ident: $ty:ty),+  $(,)?) -> $ret:ty) => {
        compile_error!("We don't support functions with more than one parameter.");
    };

    ($ocaml_name:expr, fn $name:ident($($param:ident: $ty:ty),+  $(,)?) -> $ret:ty) => {
        #[no_mangle]
        pub unsafe fn $name ($($param: $ty,)*) -> $ret {
            static FN: once_cell::sync::OnceCell<usize> = once_cell::sync::OnceCell::new();
            let the_function_to_call = *FN.get_or_init(|| {
                let the_function_to_call_name = std::ffi::CString::new($ocaml_name).expect("string contained null byte");
                let the_function_to_call = $crate::caml_named_value(the_function_to_call_name.as_ptr());
                if the_function_to_call.is_null() {
                    panic!("Could not find function. Use Callback.register");
                }
                *the_function_to_call
            });
            let args_to_function: Vec<usize> = vec![$($crate::to_ocaml(&$param),)*];
            let args_to_function_ptr: *const usize = args_to_function.as_ptr();
            let result = $crate::caml_callbackN_exn(the_function_to_call, args_to_function.len().try_into().unwrap(), args_to_function_ptr);
            if $crate::is_exception_result(result) {
                panic!("OCaml function threw an unknown exception");
            }
            let result = <$ret>::from_ocaml(result).unwrap();
            result
        }
    };

    ($ocaml_name:expr, fn $name:ident() -> $ret:ty) => {
        unsafe fn $name() -> $ret {
            $crate::ocaml_registered_function_fn!(
                $ocaml_name,
                fn inner(_unit: ()) -> $ret
            );
            inner(())
        }
    };

    ($ocaml_name:expr, fn $name:ident($($param:ident: $ty:ty),*  $(,)?)) => {
        $crate::ocaml_registered_function_fn!(
            $ocaml_name,
            fn $name($($param: $ty),*) -> ()
        );
    };
}

/// Convenience macro for declaring Rust FFI wrappers around OCaml-defined functions.
///
/// Each parameter will be converted to OCaml using `ocamlrep` and allocated on
/// the OCaml GC heap using `ocamlpool`. The result will be converted from OCaml
/// using `ocamlrep`.
///
/// Exceptions in OCaml will be caught and converted to a Rust panic. The panic
/// will not contain useful information due to the limitations of deserializing
/// arbitrary OCaml exceptions.
#[macro_export]
macro_rules! ocaml_registered_function {
    ($(fn $name:ident($($param:ident: $ty:ty),*  $(,)?) $(-> $ret:ty)?;)*) => {
        $($crate::ocaml_registered_function_fn!(
            stringify!($name),
            fn $name($($param: $ty),*) $(-> $ret)*
        );)*
    };
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

#![allow(unused_crate_dependencies)]
#![feature(exit_status_error)]

use ocamlrep_ocamlpool::ocaml_ffi;
use ocamlrep_ocamlpool::ocaml_registered_function;
use ocamlrep_ocamlpool::FromOcamlRep;

extern "C" {
    fn ocamlpool_enter();
    fn ocamlpool_reserve_block(tag: u8, size: usize) -> usize;
    fn ocamlpool_leave();
}

// This test attempts to catch off by one issues in ocamlpool.c

// Magic constant needs to fulfill two requirements:
// Needs to be above the OCAMLPOOL_DEFAULT_SIZE constant in ocamlpool.h
//   This requirement is easy to fulfill
// Needs to be the exact size of memory block allocated by ocamlpool_reserve_block
//   which is given by the Chunk_size call in chunk_alloc in ocamlpool.c
//   This requirement requires some magic
const MAGIC_MEMORY_SIZE: usize = 1053183;

ocaml_registered_function! {
    fn f_unit_to_unit();
    fn f_one_arg_to_unit(x: i64);
    fn f_sum_tuple(args: (i64, i64)) -> i64;
}

ocaml_ffi! {
    fn test() {
        unsafe {
            ocamlpool_enter();
            // This line will crash on off by one error
            ocamlpool_reserve_block(0, MAGIC_MEMORY_SIZE);
            ocamlpool_leave();
        }
    }

    fn test_call_ocaml_from_rust() {
        for _ in 0..4 {
            unsafe {
                f_unit_to_unit();
                f_one_arg_to_unit(3);
                assert!(f_sum_tuple((3, 4)) == 7);
            }
        }
    }
}

// [Note: Test blocks for Cargo]
// -----------------------------
// With buck, where testing involves compiling OCaml we make use of
// `ocaml_binary` & `custom_unnittest` rules.
//
// When testing with cargo we instead use `#[cfg(test_blocks)]` within which we
// compile OCaml "manually" using the rust `command` crate. Thus, in these cases
// the `#[cfg(test)]` blocks are for cargo only & not buck. We express that
// according to the following schema:
// ```
//  rust_library(
//     name = "foo_test",
//     ...
//     autocargo = {
//       ...
//       "test": True, // Yes, unittests for Cargo...
//       ...
//     },
//     ...
//     unittests=False, // No! No, no unittests for Buck!
//     ...
//  )
//```
//
// If in such a `#[cfg(test)]` block we now wish to use a crate not otherwise
// dependend upon & put it in the target's `deps` section in `TARGETS` there
// will be an unsused-crate error. If we put it in the target's `test_deps`
// section in `TARGETS` buck will rightly complain that `unittests=False` (so
// how can there be `tests_deps`?).
//
// The workaround I employ is to add `allow(unused_crate_dependencies)` to this
// module. That way they can be enumerated in the `deps` and are thereby
// availabe for use in the `#[cfg(test)]` blocks.

#[cfg(test)]
mod tests {
    use anyhow::Result;
    use cargo_test_utils::*;
    use tempdir::TempDir;

    #[test]
    fn ocamlpool_test() -> Result<()> {
        let tmp_dir = TempDir::new("ocamlpool_test")?;
        std::fs::copy(
            "ocamlpool_test.ml",
            tmp_dir.path().join("ocamlpool_test.ml"),
        )?;
        let compile_cmd = cmd(
            "ocamlopt.opt",
            &[
                "-verbose",
                "-c",
                "ocamlpool_test.ml",
                "-o",
                "ocamlpool_test_ml.cmx",
            ],
            Some(tmp_dir.path()),
        );
        assert_eq!(run(compile_cmd).map_err(fmt_exit_status_err), Ok(()));
        let link_cmd = cmd(
            "ocamlopt.opt",
            &[
                "-verbose",
                "-o",
                "ocamlpool_test",
                "ocamlpool_test_ml.cmx",
                "-ccopt",
                &("-L".to_owned() + workspace_dir(&["target", build_flavor()]).to_str().unwrap()),
                "-cclib",
                "-locamlpool_test",
                "-cclib",
                "-locamlrep_ocamlpool",
            ],
            Some(tmp_dir.path()),
        );
        assert_eq!(run(link_cmd).map_err(fmt_exit_status_err), Ok(()));
        let ocamlpool_test_cmd = cmd(
            tmp_dir
                .path()
                .join("ocamlpool_test")
                .as_path()
                .to_str()
                .unwrap(),
            &[],
            None,
        );
        assert_eq!(run(ocamlpool_test_cmd).map_err(fmt_exit_status_err), Ok(()));
        tmp_dir.close()?;
        Ok(())
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use syn::Meta::List;
use syn::Meta::NameValue;
use syn::Meta::Path;
use syn::NestedMeta::Lit;
use syn::NestedMeta::Meta;

static DOC: &str = "doc";
static RUST_TO_OCAML: &str = "rust_to_ocaml";
static PREFIX: &str = "prefix";
static AND: &str = "and";
static ATTR: &str = "attr";
static NAME: &str = "name";
static INLINE_TUPLE: &str = "inline_tuple";

/// The attributes understood by `rust_to_ocaml`.
#[derive(Clone, Debug)]
pub struct Attrs {
    /// Doc comments (and their desugared form, the `#[doc]` attribute) are
    /// picked up by rust_to_ocaml and included as ocamldoc comments.
    ///
    ///     /// Type A
    ///     pub type A = X;
    ///
    /// is converted to:
    ///
    ///     (** Type A *)
    ///     type a = x
    pub doc: Vec<String>,

    /// Sometimes OCaml programs use prefixes to distinguish fields of the same
    /// name in different records (avoiding OCaml warning 30). The `prefix`
    /// attribute (on a declaration of a struct or a struct-like enum variant)
    /// indicates the prefix that should be added to each field name.
    ///
    ///     #[rust_to_ocaml(prefix = "a_")]
    ///     pub struct StructA { pub foo: isize, pub bar: isize }
    ///     #[rust_to_ocaml(prefix = "b_")]
    ///     pub struct StructB { pub foo: isize, pub bar: isize }
    ///
    /// is converted to:
    ///
    ///     type struct_a = { a_foo: int; a_bar: int; }
    ///     type struct_b = { b_foo: int; b_bar: int; }
    pub prefix: Option<String>,

    /// OCaml attributes (in OCaml syntax) to annotate a type or field
    /// declaration with in the generated OCaml.
    ///
    ///     #[rust_to_ocaml(attr = "deriving show")]
    ///     pub type X = A;
    ///
    /// is converted to:
    ///
    ///     type x = a [@@deriving show]
    pub attrs: Vec<String>,

    /// Mutual recursion in type definitions is opt-in in OCaml; one writes
    /// `type x = y list and y = x list` rather than `type x = y list ;; type y
    /// = x list` (which is an error because the type name `y` is not bound in
    /// the declaration of `x`). Use the `#[rust_to_ocaml(and)]` attribute to
    /// indicate when the `and` keyword should be used to continue a mutually
    /// recursive type declaration.
    ///
    ///     pub struct Foo(pub Bar, pub Bar);
    ///     #[rust_to_ocaml(and)]
    ///     pub struct Bar(pub Option<Foo>, pub Option<Foo>);
    ///
    /// is converted to:
    ///
    ///     type foo = bar * bar
    ///     and bar = foo option * foo option
    pub mutual_rec: bool,

    /// Normally, rust_to_ocaml will convert the names of fields and enum
    /// variants by attempting to convert idiomatic Rust casing to idiomatic
    /// OCaml casing. Use the `#[rust_to_ocaml(name = "my_name")]` attribute to
    /// override this behavior and provide some other name. This attribute takes
    /// precedence over the `prefix` attribute (no prefix will be applied to the
    /// given name). This attribute cannot be used to rename types (use
    /// rust_to_ocaml_config.toml instead).
    pub name: Option<String>,

    /// In OCaml, a variant declared as `Foo of (a * b)` is a variant with one
    /// field which is a pointer to a heap-allocated tuple. A variant declared
    /// as `Baz of a * b` is a variant with two fields of type `a` and `b`.
    ///
    /// By default, rust_to_ocaml will produce variants with a single field. But
    /// this behavior can be overridden with the `inline_tuple` attribute,
    /// converting the fields of a tuple (possibly behind a reference, `Box`, or
    /// any other wrapper type declared in the `types.transparent` section in
    /// rust_to_ocaml_config.toml) to fields of the OCaml variant.
    ///
    ///     pub enum E {
    ///         Foo((A, B)),
    ///         Bar(Box<(A, B)>),
    ///         #[rust_to_ocaml(inline_tuple)]
    ///         Baz((A, B)),
    ///         #[rust_to_ocaml(inline_tuple)]
    ///         Qux(Box<(A, B)>),
    ///     }
    ///
    /// is converted to:
    ///
    ///     type e =
    ///       | Foo of (a * b)
    ///       | Bar of (a * b)
    ///       | Baz of a * b
    ///       | Qux of a * b
    pub inline_tuple: bool,
}

impl Attrs {
    pub fn from_type(item: &syn::ItemType) -> Self {
        Self::from_attributes(&item.attrs, AttrKind::Container)
    }
    pub fn from_struct(item: &syn::ItemStruct) -> Self {
        Self::from_attributes(&item.attrs, AttrKind::Container)
    }
    pub fn from_enum(item: &syn::ItemEnum) -> Self {
        Self::from_attributes(&item.attrs, AttrKind::Container)
    }

    pub fn from_variant(variant: &syn::Variant) -> Self {
        Self::from_attributes(&variant.attrs, AttrKind::Variant)
    }
    pub fn from_field(field: &syn::Field) -> Self {
        Self::from_attributes(&field.attrs, AttrKind::Field)
    }

    fn from_attributes(attrs: &[syn::Attribute], kind: AttrKind) -> Self {
        let doc = get_doc_comment(attrs);
        let mut prefix = None;
        let mut ocaml_attrs = vec![];
        let mut mutual_rec = false;
        let mut name = None;
        let mut inline_tuple = false;

        for meta_item in attrs
            .iter()
            .flat_map(get_rust_to_ocaml_meta_items)
            .flatten()
        {
            match &meta_item {
                // Parse `#[rust_to_ocaml(prefix = "foo")]`
                Meta(NameValue(m)) if m.path.is_ident(PREFIX) => {
                    // TODO: emit error for AttrKind::Field (should use the
                    // `name` meta item instead)
                    if let Ok(s) = get_lit_str(PREFIX, &m.lit) {
                        prefix = Some(s.value());
                    }
                }
                // Parse `#[rust_to_ocaml(attr = "deriving eq")]`
                Meta(NameValue(m)) if m.path.is_ident(ATTR) => {
                    if let Ok(s) = get_lit_str(ATTR, &m.lit) {
                        ocaml_attrs.push(s.value());
                    }
                }
                // Parse `#[rust_to_ocaml(and)]`
                Meta(Path(word)) if word.is_ident(AND) => {
                    // TODO: emit an error instead
                    assert_eq!(kind, AttrKind::Container);
                    mutual_rec = true;
                }
                // Parse `#[rust_to_ocaml(name = "foo")]`
                Meta(NameValue(m)) if m.path.is_ident(NAME) => {
                    // TODO: emit error for AttrKind::Container (should add to
                    // types.rename config instead)
                    if let Ok(s) = get_lit_str(NAME, &m.lit) {
                        name = Some(s.value());
                    }
                }
                // Parse `#[rust_to_ocaml(inline_tuple)]`
                Meta(Path(word)) if word.is_ident(INLINE_TUPLE) => {
                    // TODO: emit an error instead
                    assert_eq!(kind, AttrKind::Variant);
                    inline_tuple = true;
                }
                Meta(_meta_item) => {
                    // let path = meta_item
                    //     .path()
                    //     .into_token_stream()
                    //     .to_string()
                    //     .replace(' ', "");
                    // cx.error_spanned_by(
                    //     meta_item.path(),
                    //     format!("unknown rust_to_ocaml {} attribute `{}`", kind, path),
                    // );
                }
                Lit(_lit) => {
                    // cx.error_spanned_by(lit, format!("unexpected literal in rust_to_ocaml {} attribute", kind));
                }
            }
        }

        Self {
            doc,
            prefix,
            attrs: ocaml_attrs,
            mutual_rec,
            name,
            inline_tuple,
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum AttrKind {
    Container,
    Variant,
    Field,
}

impl std::fmt::Display for AttrKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Container => write!(f, "container"),
            Self::Variant => write!(f, "variant"),
            Self::Field => write!(f, "field"),
        }
    }
}

pub fn get_doc_comment(attrs: &[syn::Attribute]) -> Vec<String> {
    attrs
        .iter()
        .filter_map(|attr| {
            if !attr.path.is_ident(DOC) {
                return None;
            }
            match attr.parse_meta() {
                Ok(syn::Meta::NameValue(meta)) => {
                    if let syn::Lit::Str(s) = meta.lit {
                        Some(s.value())
                    } else {
                        None
                    }
                }
                _ => None,
            }
        })
        .collect()
}

fn get_rust_to_ocaml_meta_items(attr: &syn::Attribute) -> Result<Vec<syn::NestedMeta>, ()> {
    if !attr.path.is_ident(RUST_TO_OCAML) {
        return Ok(vec![]);
    }

    match attr.parse_meta() {
        Ok(List(meta)) => Ok(meta.nested.into_iter().collect()),
        Ok(_other) => {
            // cx.error_spanned_by(other, "expected #[rust_to_ocaml(...)]");
            Err(())
        }
        Err(_err) => {
            // cx.syn_error(err);
            Err(())
        }
    }
}

fn get_lit_str<'a>(_attr_name: &'static str, lit: &'a syn::Lit) -> Result<&'a syn::LitStr, ()> {
    if let syn::Lit::Str(lit) = lit {
        Ok(lit)
    } else {
        // cx.error_spanned_by(
        //     lit,
        //     format!(
        //         "expected rust_to_ocaml {} attribute to be a string: `{} = \"...\"`",
        //         attr_name, attr_name
        //     ),
        // );
        Err(())
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use std::str::FromStr;

use indexmap::indexmap;
use indexmap::indexset;
use indexmap::IndexMap;
use indexmap::IndexSet;
use serde::Deserialize;
use serde::Serialize;

use crate::ir;
use crate::ir::ModuleName;
use crate::ir::TypeName;

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct Config {
    #[serde(default)]
    modules: ModulesConfig,
    #[serde(default)]
    types: TypesConfig,
}

#[derive(Debug, Default, Serialize, Deserialize)]
struct ModulesConfig {
    #[serde(default, with = "indexmap::serde_seq")]
    rename: IndexMap<ModuleName, ModuleName>,
}

#[derive(Debug, Serialize, Deserialize)]
struct TypesConfig {
    transparent: IndexSet<RustTypePath>,
    #[serde(with = "indexmap::serde_seq")]
    rename: IndexMap<RustTypePath, OcamlTypePath>,
}

#[derive(Clone, PartialEq, Eq, Hash)]
pub struct RustTypePath {
    pub modules: Vec<ModuleName>,
    pub ty: TypeName,
}

#[derive(Clone, PartialEq, Eq, Hash)]
pub struct OcamlTypePath {
    pub modules: Vec<ModuleName>,
    pub ty: TypeName,
}

impl Config {
    pub fn get_renamed_module(&self, name: &ModuleName) -> Option<ModuleName> {
        self.modules.rename.get(name).cloned()
    }
    pub fn is_transparent_type(&self, path: &ir::TypePath) -> bool {
        let rust_path = RustTypePath::from(path);
        self.types.transparent.contains(&rust_path)
    }
    pub fn get_renamed_type(&self, path: &ir::TypePath) -> Option<OcamlTypePath> {
        let rust_path = RustTypePath::from(path);
        self.types.rename.get(&rust_path).cloned()
    }
}

impl Default for TypesConfig {
    fn default() -> Self {
        let r = |s| RustTypePath::from_str(s).unwrap();
        let o = |s| OcamlTypePath::from_str(s).unwrap();
        Self {
            transparent: indexset! {
                r("Box"),
                r("std::boxed::Box"),
                r("Rc"),
                r("std::rc::Rc"),
                r("Arc"),
                r("std::sync::Arc"),
            },
            rename: indexmap! {
                r("Vec") => o("list"),
                r("std::vec::Vec") => o("list"),
            },
        }
    }
}

impl From<&ir::TypePath> for RustTypePath {
    fn from(path: &ir::TypePath) -> Self {
        Self {
            modules: path.modules.clone(),
            ty: path.ty.clone(),
        }
    }
}

impl std::fmt::Display for RustTypePath {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for m in self.modules.iter() {
            write!(f, "{}", m.as_str())?;
            write!(f, "::")?;
        }
        write!(f, "{}", self.ty.as_str())
    }
}
impl std::fmt::Display for OcamlTypePath {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for m in self.modules.iter() {
            write!(f, "{}", m.as_str())?;
            write!(f, ".")?;
        }
        write!(f, "{}", self.ty.as_str())
    }
}

impl std::fmt::Debug for RustTypePath {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        format!("{self}").fmt(f)
    }
}
impl std::fmt::Debug for OcamlTypePath {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        format!("{self}").fmt(f)
    }
}

impl FromStr for RustTypePath {
    type Err = anyhow::Error;
    fn from_str(s: &str) -> anyhow::Result<Self> {
        let (modules, ty) = parse_type_path(s, "::")?;
        Ok(Self { modules, ty })
    }
}
impl FromStr for OcamlTypePath {
    type Err = anyhow::Error;
    fn from_str(s: &str) -> anyhow::Result<Self> {
        let (modules, ty) = parse_type_path(s, ".")?;
        Ok(Self { modules, ty })
    }
}

fn parse_type_path(s: &str, sep: &str) -> anyhow::Result<(Vec<ModuleName>, TypeName)> {
    let mut split = s.rsplit(sep);
    let ty = match split.next() {
        None | Some("") => anyhow::bail!("Invalid type name: {:?}", s),
        Some(ty) => TypeName(ty.to_owned()),
    };
    let mut modules = split.map(ModuleName::new).collect::<Result<Vec<_>, _>>()?;
    modules.reverse();
    Ok((modules, ty))
}

serde_from_display!(RustTypePath, "a valid Rust type path string");
serde_from_display!(OcamlTypePath, "a valid OCaml type path string");


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use anyhow::bail;
use anyhow::ensure;
use anyhow::Context;
use anyhow::Result;
use convert_case::Case;
use convert_case::Casing;

use crate::ir;
use crate::ir::Def;
use crate::ir::FieldName;
use crate::ir::File;
use crate::ir::TypeName;
use crate::ir::VariantName;
use crate::Config;

pub fn convert_file(
    config: &'static Config,
    filename: &std::path::Path,
    file: &syn::File,
) -> Result<String> {
    let defs = (file.items.iter())
        .filter_map(|item| ItemConverter::convert_item(config, item).transpose())
        .collect::<Result<_>>()?;
    let file_stem = filename.file_stem().context("expected nonempty filename")?;
    let module_name = file_stem.to_str().context("non-UTF8 filename")?.to_owned();
    let mut file = File {
        root: ir::Module {
            name: ir::ModuleName::new(module_name)?,
            defs,
        },
    };
    crate::rewrite_types::rewrite_file(config, &mut file);
    crate::rewrite_module_names::rewrite_file(config, &mut file);
    Ok(file.to_string())
}

struct ItemConverter {
    config: &'static Config,
    tparams: Vec<String>,
}

impl ItemConverter {
    fn convert_item(config: &'static Config, item: &syn::Item) -> Result<Option<Def>> {
        use syn::Item;
        match item {
            Item::Type(item) => {
                let this = ItemConverter::new(config, &item.generics);
                Ok(Some(this.convert_item_type(item).with_context(|| {
                    format!("Failed to convert type {}", item.ident)
                })?))
            }
            Item::Struct(item) => {
                let this = ItemConverter::new(config, &item.generics);
                Ok(Some(this.convert_item_struct(item).with_context(|| {
                    format!("Failed to convert type {}", item.ident)
                })?))
            }
            Item::Enum(item) => {
                let this = ItemConverter::new(config, &item.generics);
                Ok(Some(this.convert_item_enum(item).with_context(|| {
                    format!("Failed to convert type {}", item.ident)
                })?))
            }
            Item::Mod(item) => {
                if let Some((_brace, items)) = &item.content {
                    let defs = items
                        .iter()
                        .filter_map(|item| Self::convert_item(config, item).transpose())
                        .collect::<Result<_>>()
                        .with_context(|| format!("Failed to convert module {}", item.ident))?;
                    Ok(Some(Def::Module(ir::Module {
                        name: ir::ModuleName::new(item.ident.to_string())?,
                        defs,
                    })))
                } else {
                    Ok(None)
                }
            }
            _ => Ok(None),
        }
    }

    fn new(config: &'static Config, generics: &syn::Generics) -> Self {
        let tparams = generics
            .type_params()
            .map(|tparam| tparam.ident.to_string())
            .collect();
        Self { config, tparams }
    }

    fn convert_item_type(self, item: &syn::ItemType) -> Result<Def> {
        let name = TypeName(item.ident.to_string().to_case(Case::Snake));
        let attrs = attr_parser::Attrs::from_type(item);
        let ty = self.convert_type(&item.ty)?;
        Ok(Def::Alias {
            doc: attrs.doc,
            attrs: attrs.attrs,
            mutual_rec: attrs.mutual_rec,
            tparams: self.tparams,
            name,
            ty,
        })
    }

    fn convert_item_struct(self, item: &syn::ItemStruct) -> Result<Def> {
        let name = TypeName(item.ident.to_string().to_case(Case::Snake));
        let container_attrs = attr_parser::Attrs::from_struct(item);
        match &item.fields {
            syn::Fields::Unit => Ok(Def::Alias {
                doc: container_attrs.doc,
                attrs: container_attrs.attrs,
                mutual_rec: container_attrs.mutual_rec,
                tparams: self.tparams,
                name,
                ty: ir::Type::Path(ir::TypePath::simple("unit")),
            }),
            syn::Fields::Unnamed(fields) => {
                let elems = (fields.unnamed.iter())
                    .map(|field| self.convert_type(&field.ty))
                    .collect::<Result<Vec<_>>>()?;
                Ok(Def::Alias {
                    doc: container_attrs.doc,
                    attrs: container_attrs.attrs,
                    mutual_rec: container_attrs.mutual_rec,
                    tparams: self.tparams,
                    name,
                    ty: if elems.is_empty() {
                        ir::Type::Path(ir::TypePath::simple("unit"))
                    } else {
                        ir::Type::Tuple(ir::TypeTuple { elems })
                    },
                })
            }
            syn::Fields::Named(fields) => {
                let fields = (fields.named.iter())
                    .map(|field| {
                        let field_attrs = attr_parser::Attrs::from_field(field);
                        let name = if let Some(name) = field_attrs.name {
                            FieldName(name)
                        } else {
                            field_name(field.ident.as_ref(), container_attrs.prefix.as_deref())
                        };
                        let ty = self.convert_type(&field.ty)?;
                        Ok(ir::Field {
                            name,
                            ty,
                            doc: field_attrs.doc,
                            attrs: field_attrs.attrs,
                        })
                    })
                    .collect::<Result<Vec<_>>>()?;
                Ok(Def::Record {
                    doc: container_attrs.doc,
                    attrs: container_attrs.attrs,
                    mutual_rec: container_attrs.mutual_rec,
                    tparams: self.tparams,
                    name,
                    fields,
                })
            }
        }
    }

    fn convert_item_enum(self, item: &syn::ItemEnum) -> Result<Def> {
        let name = TypeName(item.ident.to_string().to_case(Case::Snake));
        let container_attrs = attr_parser::Attrs::from_enum(item);
        let variants = item
            .variants
            .iter()
            .map(|variant| {
                let variant_attrs = attr_parser::Attrs::from_variant(variant);
                let name = if let Some(name) = variant_attrs.name {
                    VariantName(name)
                } else {
                    variant_name(&variant.ident, container_attrs.prefix.as_deref())
                };
                let fields = match &variant.fields {
                    syn::Fields::Unit => None,
                    syn::Fields::Unnamed(fields) => {
                        let mut fields = (fields.unnamed.iter())
                            .map(|field| self.convert_type(&field.ty))
                            .collect::<Result<Vec<_>>>()?;
                        if variant_attrs.inline_tuple {
                            assert_eq!(fields.len(), 1);
                            let field = fields.pop().unwrap();
                            match field {
                                ir::Type::Path(mut path) => {
                                    if path.targs.len() == 1
                                        && self.config.is_transparent_type(&path)
                                        && matches!(path.targs[0], ir::Type::Tuple(..))
                                    {
                                        if let Some(ir::Type::Tuple(tuple)) = path.targs.pop() {
                                            Some(ir::VariantFields::Unnamed(tuple.elems))
                                        } else {
                                            unreachable!()
                                        }
                                    } else {
                                        anyhow::bail!(
                                            "Variant {} must have a single argument which is a tuple",
                                            variant.ident.to_string()
                                        )
                                    }
                                }
                                ir::Type::Tuple(tuple) => {
                                    Some(ir::VariantFields::Unnamed(tuple.elems))
                                }
                            }
                        } else {
                            Some(ir::VariantFields::Unnamed(fields))
                        }
                    }
                    syn::Fields::Named(fields) => Some(ir::VariantFields::Named(
                        (fields.named.iter())
                            .map(|field| {
                                let field_attrs = attr_parser::Attrs::from_field(field);
                                let name = if let Some(name) = field_attrs.name {
                                    FieldName(name)
                                } else {
                                    field_name(
                                        field.ident.as_ref(),
                                        variant_attrs.prefix.as_deref(),
                                    )
                                };
                                let ty = self.convert_type(&field.ty)?;
                                Ok(ir::Field {
                                    name,
                                    ty,
                                    doc: field_attrs.doc,
                                    attrs: field_attrs.attrs,
                                })
                            })
                            .collect::<Result<_>>()?,
                    )),
                };
                Ok(ir::Variant {
                    name,
                    fields,
                    doc: variant_attrs.doc,
                    attrs: variant_attrs.attrs,
                })
            })
            .collect::<Result<Vec<_>>>()?;
        Ok(Def::Variant {
            doc: container_attrs.doc,
            attrs: container_attrs.attrs,
            mutual_rec: container_attrs.mutual_rec,
            tparams: self.tparams,
            name,
            variants,
        })
    }

    fn convert_type(&self, ty: &syn::Type) -> Result<ir::Type> {
        match ty {
            syn::Type::Path(ty) => Ok(ir::Type::Path(self.convert_type_path(ty)?)),
            syn::Type::Tuple(ty) => {
                if ty.elems.is_empty() {
                    Ok(ir::Type::Path(ir::TypePath::simple("unit")))
                } else {
                    Ok(ir::Type::Tuple(ir::TypeTuple {
                        elems: (ty.elems.iter())
                            .map(|e| self.convert_type(e))
                            .collect::<Result<_>>()?,
                    }))
                }
            }
            syn::Type::Reference(ty) => Ok(self.convert_type(&ty.elem)?),
            syn::Type::Slice(ty) => Ok(ir::Type::Path(ir::TypePath {
                modules: vec![],
                targs: vec![self.convert_type(&ty.elem)?],
                ty: ir::TypeName(String::from("list")),
            })),
            _ => bail!("Not supported: {:?}", ty),
        }
    }

    fn convert_type_path(&self, ty: &syn::TypePath) -> Result<ir::TypePath> {
        ensure!(ty.qself.is_none(), "Qualified self in paths not supported");
        let last_seg = ty.path.segments.last().unwrap();
        if ty.path.segments.len() == 1 && last_seg.arguments.is_empty() {
            let ident = last_seg.ident.to_string();
            if self.tparams.contains(&ident) {
                let tparam = format!("'{}", ident.to_case(Case::Snake));
                return Ok(ir::TypePath::simple(tparam));
            }
        }
        let segments_except_last = ty.path.segments.iter().rev().skip(1).rev();
        Ok(ir::TypePath {
            modules: segments_except_last
                .map(|seg| {
                    ensure!(
                        seg.arguments.is_empty(),
                        "Type args only supported in last path segment"
                    );
                    ir::ModuleName::new(seg.ident.to_string())
                })
                .collect::<Result<_>>()?,
            ty: TypeName(last_seg.ident.to_string()),
            targs: match &last_seg.arguments {
                syn::PathArguments::AngleBracketed(args) => (args.args.iter())
                    .filter_map(|arg| match arg {
                        syn::GenericArgument::Type(arg) => Some(self.convert_type(arg)),
                        _ => None,
                    })
                    .collect::<Result<_>>()?,
                _ => vec![],
            },
        })
    }
}

fn field_name(ident: Option<&syn::Ident>, prefix: Option<&str>) -> FieldName {
    FieldName(format!("{}{}", prefix.unwrap_or_default(), ident.unwrap()))
}

fn variant_name(ident: &syn::Ident, prefix: Option<&str>) -> VariantName {
    VariantName(format!("{}{}", prefix.unwrap_or_default(), ident))
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

mod display;

use derive_more::Display;

#[derive(Debug)]
pub struct File {
    pub root: Module,
}

#[derive(Debug)]
pub struct Module {
    pub name: ModuleName,
    pub defs: Vec<Def>,
}

#[derive(Debug)]
pub enum Def {
    Module(Module),
    Alias {
        doc: Vec<String>,
        attrs: Vec<String>,
        mutual_rec: bool,
        tparams: Vec<String>,
        name: TypeName,
        ty: Type,
    },
    Record {
        doc: Vec<String>,
        attrs: Vec<String>,
        mutual_rec: bool,
        tparams: Vec<String>,
        name: TypeName,
        fields: Vec<Field>,
    },
    Variant {
        doc: Vec<String>,
        attrs: Vec<String>,
        mutual_rec: bool,
        tparams: Vec<String>,
        name: TypeName,
        variants: Vec<Variant>,
    },
}

#[derive(Debug)]
pub struct Variant {
    pub name: VariantName,
    pub fields: Option<VariantFields>,
    pub doc: Vec<String>,
    pub attrs: Vec<String>,
}

#[derive(Debug)]
pub enum VariantFields {
    Unnamed(Vec<Type>),
    Named(Vec<Field>),
}

#[derive(Debug)]
pub struct Field {
    pub name: FieldName,
    pub ty: Type,
    pub doc: Vec<String>,
    pub attrs: Vec<String>,
}

#[derive(Debug)]
pub enum Type {
    Path(TypePath),
    Tuple(TypeTuple),
}

#[derive(Debug)]
pub struct TypePath {
    pub targs: Vec<Type>,
    pub modules: Vec<ModuleName>,
    pub ty: TypeName,
}

impl TypePath {
    pub fn simple(id: impl Into<String>) -> Self {
        Self {
            modules: vec![],
            ty: TypeName(id.into()),
            targs: vec![],
        }
    }
}

#[derive(Debug)]
pub struct TypeTuple {
    pub elems: Vec<Type>,
}

#[derive(Clone, Hash, PartialEq, Eq)]
pub struct ModuleName(String);

impl ModuleName {
    pub fn new(name: impl Into<String>) -> anyhow::Result<Self> {
        let name: String = name.into();
        anyhow::ensure!(!name.is_empty(), "Module names must not be empty");
        let first_char = name.chars().next().unwrap();
        anyhow::ensure!(
            first_char.is_ascii(),
            "Module names must start with an ASCII character: {}",
            name
        );
        anyhow::ensure!(
            first_char.to_ascii_uppercase().is_ascii_uppercase(),
            "Module names must start with a character which can be converted to uppercase: {}",
            name
        );
        Ok(Self(name))
    }

    pub fn as_str(&self) -> &str {
        self.0.as_str()
    }
}

impl std::str::FromStr for ModuleName {
    type Err = anyhow::Error;
    fn from_str(s: &str) -> anyhow::Result<Self> {
        Self::new(s)
    }
}

serde_from_display!(ModuleName, "a valid Rust or OCaml module name");

#[derive(Clone, Hash, PartialEq, Eq)]
pub struct TypeName(pub String);

impl TypeName {
    pub fn as_str(&self) -> &str {
        self.0.as_str()
    }
}

#[derive(Clone, Hash, PartialEq, Eq)]
pub struct FieldName(pub String);

#[derive(Clone, Hash, PartialEq, Eq, Display)]
pub struct VariantName(pub String);

impl std::fmt::Debug for ModuleName {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.fmt(f)
    }
}
impl std::fmt::Debug for TypeName {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.fmt(f)
    }
}
impl std::fmt::Debug for FieldName {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.fmt(f)
    }
}
impl std::fmt::Debug for VariantName {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.fmt(f)
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

/// Provide impls of `Serialize` and `Deserialize` which delegate to the impls
/// of `std::fmt::Display` and `std::str::FromStr` respectively.
macro_rules! serde_from_display {
    ($name:ident, $expecting:expr) => {
        impl ::serde::Serialize for $name {
            fn serialize<S: ::serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
                serializer.serialize_str(&self.to_string())
            }
        }
        impl<'de> ::serde::Deserialize<'de> for $name {
            fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
            where
                D: ::serde::Deserializer<'de>,
            {
                struct Visitor;
                impl<'de> ::serde::de::Visitor<'de> for Visitor {
                    type Value = $name;
                    fn expecting(
                        &self,
                        formatter: &mut ::std::fmt::Formatter<'_>,
                    ) -> ::std::fmt::Result {
                        formatter.write_str($expecting)
                    }
                    fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
                    where
                        E: ::serde::de::Error,
                    {
                        value.parse().map_err(|e| {
                            E::invalid_value(::serde::de::Unexpected::Other(&format!("{e}")), &self)
                        })
                    }
                }
                deserializer.deserialize_str(Visitor)
            }
        }
    };
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use crate::ir;
use crate::Config;

pub fn rewrite_file(config: &'static Config, file: &mut ir::File) {
    let rewriter = Rewriter { config };
    rewriter.rewrite_module(&mut file.root)
}

struct Rewriter {
    config: &'static Config,
}

impl Rewriter {
    fn rewrite_module(&self, module: &mut ir::Module) {
        module.defs.iter_mut().for_each(|def| self.rewrite_def(def))
    }

    fn rewrite_def(&self, def: &mut ir::Def) {
        match def {
            ir::Def::Module(module) => self.rewrite_module(module),
            ir::Def::Alias { ty, .. } => self.rewrite_type(ty),
            ir::Def::Record { fields, .. } => fields.iter_mut().for_each(|f| self.rewrite_field(f)),
            ir::Def::Variant { variants, .. } => {
                variants.iter_mut().for_each(|v| self.rewrite_variant(v))
            }
        }
    }

    fn rewrite_field(&self, field: &mut ir::Field) {
        self.rewrite_type(&mut field.ty)
    }

    fn rewrite_variant(&self, variant: &mut ir::Variant) {
        variant
            .fields
            .iter_mut()
            .for_each(|f| self.rewrite_variant_fields(f))
    }

    fn rewrite_variant_fields(&self, fields: &mut ir::VariantFields) {
        match fields {
            ir::VariantFields::Unnamed(tys) => tys.iter_mut().for_each(|ty| self.rewrite_type(ty)),
            ir::VariantFields::Named(fields) => {
                fields.iter_mut().for_each(|f| self.rewrite_field(f))
            }
        }
    }

    fn rewrite_type(&self, ty: &mut ir::Type) {
        match ty {
            ir::Type::Path(path) => self.rewrite_type_path(path),
            ir::Type::Tuple(tuple) => self.rewrite_type_tuple(tuple),
        }
    }

    fn rewrite_type_path(&self, path: &mut ir::TypePath) {
        match path.modules.get(0).map(ir::ModuleName::as_str) {
            Some("crate" | "super") => {
                path.modules.remove(0);
            }
            _ => {}
        }
        path.modules.iter_mut().for_each(|m| {
            if let Some(name) = self.config.get_renamed_module(m) {
                *m = name;
            }
        });
        path.targs.iter_mut().for_each(|ty| self.rewrite_type(ty))
    }

    fn rewrite_type_tuple(&self, tuple: &mut ir::TypeTuple) {
        tuple.elems.iter_mut().for_each(|ty| self.rewrite_type(ty))
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

use convert_case::Case;
use convert_case::Casing;

use crate::ir;
use crate::Config;

pub fn rewrite_file(config: &'static Config, file: &mut ir::File) {
    Rewriter::rewrite_module(config, &mut file.root)
}

struct Rewriter {
    config: &'static Config,
    module_name: ir::ModuleName,
}

impl Rewriter {
    fn rewrite_module(config: &'static Config, module: &mut ir::Module) {
        let this = Self {
            config,
            module_name: module.name.clone(),
        };
        module.defs.iter_mut().for_each(|def| this.rewrite_def(def))
    }

    fn rewrite_def(&self, def: &mut ir::Def) {
        let rewrite_name = |name: &mut ir::TypeName| {
            if name.as_str() == self.module_name.as_str() {
                *name = ir::TypeName(String::from("t"));
            }
        };
        match def {
            ir::Def::Module(module) => Self::rewrite_module(self.config, module),
            ir::Def::Alias { name, ty, .. } => {
                rewrite_name(name);
                self.rewrite_type(ty)
            }
            ir::Def::Record { name, fields, .. } => {
                rewrite_name(name);
                fields.iter_mut().for_each(|f| self.rewrite_field(f))
            }
            ir::Def::Variant { name, variants, .. } => {
                rewrite_name(name);
                variants.iter_mut().for_each(|v| self.rewrite_variant(v))
            }
        }
    }

    fn rewrite_field(&self, field: &mut ir::Field) {
        self.rewrite_type(&mut field.ty)
    }

    fn rewrite_variant(&self, variant: &mut ir::Variant) {
        variant
            .fields
            .iter_mut()
            .for_each(|f| self.rewrite_variant_fields(f))
    }

    fn rewrite_variant_fields(&self, fields: &mut ir::VariantFields) {
        match fields {
            ir::VariantFields::Unnamed(tys) => tys.iter_mut().for_each(|ty| self.rewrite_type(ty)),
            ir::VariantFields::Named(fields) => {
                fields.iter_mut().for_each(|f| self.rewrite_field(f))
            }
        }
    }

    fn rewrite_type(&self, ty: &mut ir::Type) {
        match ty {
            ir::Type::Path(path) => {
                if path.targs.len() == 1 && self.config.is_transparent_type(path) {
                    *ty = path.targs.pop().unwrap();
                    self.rewrite_type(ty);
                } else {
                    self.rewrite_type_path(path);
                }
            }
            ir::Type::Tuple(tuple) => self.rewrite_type_tuple(tuple),
        }
    }

    fn rewrite_type_path(&self, path: &mut ir::TypePath) {
        // Convert all integer types to `int`. The impls of ToOcamlRep
        // and FromOcamlRep for integer types do checked conversions, so
        // we'll fail at runtime if our int value doesn't fit into
        // OCaml's integer width.
        if path.modules.is_empty() && path.targs.is_empty() {
            match path.ty.as_str() {
                "i8" | "u8" | "i16" | "u16" | "i32" | "u32" | "i64" | "u64" | "i128" | "u128"
                | "isize" | "usize" => {
                    path.ty = ir::TypeName(String::from("int"));
                    return;
                }
                _ => {}
            }
        }
        if let Some(renamed_path) = self.config.get_renamed_type(path) {
            path.ty = renamed_path.ty;
            path.modules = renamed_path.modules;
        }
        let ty = path.ty.as_str().to_case(Case::Snake);
        let ty_matches_last_module_in_path =
            (path.modules.last()).map_or(false, |module| ty == module.as_str());
        if ty_matches_last_module_in_path || ty == self.module_name.as_str() {
            path.ty = ir::TypeName(String::from("t"));
        }
        path.targs.iter_mut().for_each(|ty| self.rewrite_type(ty))
    }

    fn rewrite_type_tuple(&self, tuple: &mut ir::TypeTuple) {
        tuple.elems.iter_mut().for_each(|ty| self.rewrite_type(ty))
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

#[macro_use]
mod macros;

mod config;
mod convert;
mod ir;
mod rewrite_module_names;
mod rewrite_types;

use std::io::Write;
use std::path::Path;
use std::path::PathBuf;

use anyhow::Context;
use anyhow::Result;

use crate::config::Config;

#[derive(Debug, clap::Parser)]
struct Opts {
    /// The Hack source file to convert.
    #[clap(value_name("FILEPATH"))]
    filename: PathBuf,

    /// The OCaml source file to generate.
    #[clap(value_name("OUTPATH"))]
    out_path: Option<PathBuf>,

    /// Path to a configuration file.
    #[clap(long)]
    config: Option<PathBuf>,

    /// Command to regenerate the output. This text will be included in generated file headers.
    #[clap(long)]
    regen_cmd: Option<String>,

    /// Do not add copyright header and generated tag (for tests).
    #[clap(long)]
    no_header: bool,

    /// Path to an OCaml formatter binary which will be used on the generated output.
    #[clap(long)]
    formatter: Option<String>,
}

fn main() -> Result<()> {
    let opts = <Opts as clap::Parser>::from_args();

    let config = Box::leak(Box::new(match opts.config {
        Some(path) => {
            let contents = std::fs::read_to_string(&path)
                .with_context(|| format!("Failed to read config file at {}", path.display()))?;
            toml::from_str(&contents)
                .with_context(|| format!("Failed to parse config file at {}", path.display()))?
        }
        None => Config::default(),
    }));

    let src = std::fs::read_to_string(&opts.filename)
        .with_context(|| format!("Failed to read input file {}", opts.filename.display()))?;
    let file = syn::parse_file(&src)?;
    let mut ocaml_src = convert::convert_file(config, &opts.filename, &file)?;

    if !opts.no_header {
        ocaml_src = attach_header(opts.regen_cmd.as_deref(), &ocaml_src);
    }
    let absolute_filename = opts.filename.canonicalize()?;
    let mut ocaml_src = ocamlformat(
        opts.formatter.as_deref(),
        opts.out_path
            .as_deref()
            .and_then(Path::parent)
            .or_else(|| absolute_filename.parent()),
        ocaml_src.into_bytes(),
    )
    .context("failed to run ocamlformat")?;
    if !opts.no_header {
        ocaml_src = signed_source::sign_file(&ocaml_src)?;
    }

    if let Some(out_path) = &opts.out_path {
        write_file(out_path, &ocaml_src)?;
    } else {
        let mut stdout = std::io::stdout().lock();
        stdout.write_all(&ocaml_src)?;
    }

    Ok(())
}

fn attach_header(regen_cmd: Option<&str>, contents: &str) -> String {
    let regen_cmd = regen_cmd.map_or_else(String::new, |regen_cmd| {
        format!(" *\n * To regenerate this file, run:\n *   {}\n", regen_cmd)
    });
    format!(
        r#"(*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *
 * {}
{} *)

{}"#,
        signed_source::SIGNING_TOKEN,
        regen_cmd,
        contents
    )
}

fn ocamlformat(
    formatter: Option<&str>,
    out_dir: Option<&Path>,
    contents: Vec<u8>,
) -> Result<Vec<u8>> {
    let formatter = match formatter {
        None => return Ok(contents),
        Some(f) => f,
    };
    // Even if we format the file on disk (i.e., at `opts.out_path`),
    // ocamlformat won't look for an .ocamlformat file in the directory
    // containing the file. It only looks up from the current working directory.
    // There's a --root arg, but it doesn't seem to produce the same behavior.
    let prev_dir = std::env::current_dir()?;
    if let Some(out_dir) = out_dir {
        std::env::set_current_dir(out_dir)?;
    }
    let mut child = std::process::Command::new(formatter)
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        // In the event that an .ocamlformat file is still not available, tell
        // ocamlformat to please format it anyway.
        .arg("--enable-outside-detected-project")
        .arg("--impl")
        .arg("-")
        .spawn()?;
    let child_stdin = child.stdin.as_mut().unwrap();
    child_stdin.write_all(&contents)?;
    let output = child.wait_with_output()?;
    if !output.status.success() {
        anyhow::bail!("Formatter failed:\n{:#?}", output);
    }
    if out_dir.is_some() {
        std::env::set_current_dir(prev_dir)?;
    }
    Ok(output.stdout)
}

fn write_file(path: &Path, contents: &[u8]) -> Result<()> {
    let mut file = std::fs::File::create(path)?;
    file.write_all(contents)?;
    Ok(())
}