< MizarさんのAI生成コード置き場

fastout

投稿者Mizar
言語Rust
更新日時 (JST)2026-10-17 18:19:28
サイズ23,094 Byte
参照URL

日時はサーバーが記録したものです。更新すると更新日時が変わるので、不都合がある場合は新しく投稿してください。

テキストで開く
// ============================================================================
// FastFormat: 10^19 高速除算とゼロ埋め動的長出力を駆使した、u128 までの整数型に対する超高速フォーマッタ
// ============================================================================
pub use io::Write;
use std::{
    ptr::{copy_nonoverlapping as c, write_unaligned as w},
    *,
};
type R = io::Result<()>;

const D8: u64 = 10u64.pow(8); // 10^8
const D16: u64 = 10u64.pow(16); // 10^16
const D32: u128 = 10u128.pow(32); // 10^32
const D32U: u64 = (D32 >> 64) as u64; // 10^32 の上位64bit

// const D19: u64 = 10u64.pow(19); // 10^19

// ============================================================================
// デバッグビルドでは通常の assert! と同様に動作し、リリースビルドでは条件が偽のときに UB を発生させる関数
// (リリースビルドで条件が偽のときに UB を発生させることで、コンパイラが条件を常に真と仮定して最適化できるようになる)
// ============================================================================
#[inline(always)]
pub unsafe fn assert_unchecked_pseudo(cond: bool) {
    if cfg!(debug_assertions) {
        assert!(cond);
    } else {
        if !cond {
            // cond が偽でここに来た時点で UB(= assert_unchecked と同じ性質)
            std::hint::unreachable_unchecked();
        }
    }
}

// ============================================================================
// 10^4 のゼロ埋め固定長のLUT技法 (v < 10000 の前提)
// 0〜9999 の全パターンを事前にアスキー化して LUT に格納
// ============================================================================
static LP: [u32; 10000] = {
    let (mut x, mut i) = ([0u32; 10000], 0);
    while i < 10000 {
        let (a, b, c, d) = (i / 1000, i / 100 % 10, i / 10 % 10, i % 10);
        x[i as usize] = 0x30303030 | a | (b << 8) | (c << 16) | (d << 24);
        i += 1;
    }
    x
};

// ============================================================================
// 2桁のゼロ埋め動的長出力のための、0〜99 の全パターンをアスキー化して LUT に格納
// 例: 45 -> 0x0504 ('4' = 0x34, '5' = 0x35, パッキングして 0x3534 -> 0x0504)
// ============================================================================
static LH: [u16; 100] = {
    let (mut x, mut i) = ([0u16; 100], 0);
    while i < 100 {
        let (h, l) = (i / 10, i % 10);
        x[i as usize] = 0x3030 | h | (l << 8);
        i += 1;
    }
    x
};

// ============================================================================
// 汎用の商と余りを同時に計算する関数 (商 q と余り r を返す)
// ============================================================================
#[inline(always)]
fn dr<U>(a: U, b: U) -> (U, U)
where
    U: ops::Div<Output = U> + ops::Rem<Output = U> + Copy,
{
    (a / b, a % b)
}

// ============================================================================
// 入力 x は 0 <= x < 10^32 を前提とし、16桁ごとに分割した2つの整数を返す (x = q * 10^16 + r)
// ============================================================================
#[inline(always)]
fn div16l(x: u128) -> (u64, u64) {
    // floor(2^116 / 10^16) = 8307674973655724205
    let mut s = (((x >> 52) as u64 as u128 * 8307674973655724205) >> 64) as u64;
    let mut t = (x as u64).wrapping_sub(s.wrapping_mul(D16));
    let d = t >= D16;
    s += d as u64;
    t -= D16 * d as u64;
    (s, t)
}

// ============================================================================
// 入力 x は 0 <= x < 2^128 を前提とし、
// 16桁ごとに分割した3つの整数を返す (x = q * 10^32 + r * 10^16 + s)
// ============================================================================
#[inline(always)]
fn div16(x: u128) -> (u32, u64, u64) {
    let (mut q, mut t) = (0, x);
    if ((x >> 64) as u64) >= D32U {
        // floor(2^128 / 10^32) = 3402823
        q = (((x >> 64) as u64 as u128 * 3402823) >> 64) as u32;
        t = x - (q as u128 * D32);
        let d = t >= D32;
        q += d as u32;
        t -= D32 * d as u128;
    }
    let (r, s) = div16l(t);
    (q, r, s)
}

/*
// ============================================================================
// Lean 4 で証明された 127bit用 10^19 高速除算モジュール
// 入力 x は 0 <= x < 2^127 を前提とし、商 q と余り r を返す
// ============================================================================
#[inline(always)]
fn div19(x: u128) -> (u64, u64) {
    // x を 2^63 で右シフトして、上位ビットを抽出 (127bit の場合は最大 2^127-1 なので、シフト後は最大 2^64-1 になる)
    // 乗算と上位抽出は u64 × u64 -> 128bit のハードウェア命令(mul)に乗せる
    // floor(2^64 / 10^19) = 17014118346046923173
    let mut q = (((x >> 63) as u64 as u128 * 17014118346046923173) >> 64) as u64;

    // 余りの計算と判定は、必ず u128 の空間で行う!
    // (x は最大 2^127-1、q * D も同スケールなので引き算自体は u128 で安全)
    let mut r = x - (q as u128 * D19 as u128);

    // 最大誤差2の補正
    let d = r >= D19 as u128;
    q += d as u64;
    r -= D19 as u128 * d as u128;
    let d = r >= D19 as u128;
    q += d as u64;
    r -= D19 as u128 * d as u128;

    // 補正完了後、r は確実に 10^19 未満 (u64::MAX 以下) になるため、ここで初めて安全にキャスト
    (q, r as u64)
}
*/

/*
#[inline(always)]
#[cfg(target_arch = "x86_64")]
unsafe fn div19(x: u128) -> (u64, u64) {
    assert_unchecked_pseudo((x >> 64) < D19 as u128); // x < 2^64 * 10^19 を前提とする
    unsafe {
        let (y0, y1);
        arch::asm!(
            "div {0}",
            in(reg) 10u64.pow(19),
            inlateout("rax") x as u64 => y0,
            inlateout("rdx") (x >> 64) as u64 => y1,
        );
        (y0, y1)
    }
}
*/

// ============================================================================
// 2桁のゼロ埋め固定長 (v < 100 の前提)
// ============================================================================
#[inline(always)]
unsafe fn l2(v: u32) -> u16 {
    unsafe { assert_unchecked_pseudo(v < 100) };
    LH[v as usize]
}

// ============================================================================
// 4桁のゼロ埋め固定長 (v < 10000 の前提)
// ============================================================================
#[inline(always)]
unsafe fn l4(v: u32) -> u32 {
    unsafe { assert_unchecked_pseudo(v < 10000) };
    LP[v as usize]
}

// ============================================================================
// 8桁のゼロ埋め固定長 (v < 100 million の前提)
// ============================================================================
#[inline(always)]
unsafe fn l8(v: u32) -> u64 {
    let (h, l) = dr(v, 10000);
    unsafe { (l4(h) as u64) | ((l4(l) as u64) << 32) }
}

// ============================================================================
// 8桁のゼロ埋め固定長 (v < 100 million の前提)
// 常に8文字分のスペースがバッファに書き込まれる
// ============================================================================
#[inline(always)]
unsafe fn f8(p: *mut u8, v: u32) {
    unsafe { w(p as *mut u64, l8(v).to_le()) };
}

// ============================================================================
// 16桁のゼロ埋め固定長 (v < 10^16 の前提)
// 常に16文字分のスペースがバッファに書き込まれる
// ============================================================================
#[inline(always)]
unsafe fn f16(p: *mut u8, v: u64) {
    let (h, l) = dr(v, D8);
    unsafe {
        f8(p, h as u32);
        f8(p.add(8), l as u32);
    }
}

/*
// ============================================================================
// 19桁のゼロ埋め固定長 (v < 10^19 の前提)
// 常に19文字分のスペースがバッファに書き込まれる
// ============================================================================
#[inline(always)]
unsafe fn f19(p: *mut u8, v: u64) {
    let (h, l) = dr(v, D16);
    unsafe {
        w(p as *mut u32, (l4(h as u32) >> 8).to_le());
        f16(p.add(3), l);
    }
}
*/

// ============================================================================
// 2桁のゼロ埋め動的長 (v < 100 の前提)
// ============================================================================
#[inline(always)]
unsafe fn d2(p: *mut u8, v: u32) -> usize {
    unsafe {
        let a = l2(v);
        let s = (a & 0xff0f).trailing_zeros();
        w(p as *mut u16, (a >> (s & !7)).to_le());
        2 - (s / 8) as usize
    }
}

// ============================================================================
// 4桁のゼロ埋め動的長 (v < 10000 の前提)
// ============================================================================
#[inline(always)]
unsafe fn d4(p: *mut u8, v: u32) -> usize {
    unsafe {
        let a = l4(v);
        let s = (a & 0xff0f0f0f).trailing_zeros();
        w(p as *mut u32, (a >> (s & !7)).to_le());
        4 - (s / 8) as usize
    }
}

// ============================================================================
// 8桁のゼロ埋め動的長 (v < 100 million の前提)
// ============================================================================
#[inline(always)]
unsafe fn d8(p: *mut u8, v: u32) -> usize {
    unsafe {
        let a = l8(v);
        let s = (a & 0xff0f0f0f0f0f0f0f).trailing_zeros();
        w(p as *mut u64, (a >> (s & !7)).to_le());
        8 - (s / 8) as usize
    }
}

// ============================================================================
// 16桁のゼロ埋め動的長 (v < 10^16 の前提)
// ============================================================================
#[inline(always)]
unsafe fn d16(p: *mut u8, v: u64) -> usize {
    unsafe {
        if v < D8 {
            d8(p, v as u32)
        } else {
            let (h, l) = dr(v, D8);
            let o = d8(p, h as u32);
            f8(p.add(o), l as u32);
            o + 8
        }
    }
}

// ============================================================================
// 符号なし 32bit フォーマッタ
// ============================================================================
#[inline(always)]
unsafe fn du32(p: *mut u8, v: u32) -> usize {
    unsafe {
        if v < 100000000 {
            d8(p, v)
        } else {
            let (h, l) = dr(v, 100000000);
            let o = d2(p, h);
            f8(p.add(o), l);
            o + 8
        }
    }
}

// ============================================================================
// 符号なし 64bit フォーマッタ
// ============================================================================
unsafe fn du64(p: *mut u8, v: u64) -> usize {
    unsafe {
        if v < D16 {
            d16(p, v)
        } else {
            let (h, l) = dr(v, D16);
            let o = d4(p, h as u32);
            f16(p.add(o), l);
            o + 16
        }
    }
}

// ============================================================================
// 符号なし 128bit フォーマッタ
// ============================================================================
#[inline(always)]
unsafe fn du128(p: *mut u8, v: u128) -> usize {
    unsafe {
        if (v >> 64) == 0 {
            du64(p, v as u64)
        } else {
            let (h, m, l) = div16(v);
            if h == 0 {
                let o = d16(p, m);
                f16(p.add(o), l);
                o + 16
            } else {
                let o = d8(p, h as u32);
                f16(p.add(o), m);
                f16(p.add(o + 16), l);
                o + 32
            }
        }
    }
}

// ============================================================================
// 符号付き整数型のフォーマッタは、符号なし版を呼び出すだけの薄いラッパーで実装
// ============================================================================
macro_rules! impl_ifmt {
    ($i:ty, $u:ty, $f:ident, $g:ident) => {
        #[inline(always)]
        unsafe fn $f(ptr: *mut u8, v: $i) -> usize {
            // v < 0 なら 1、v >= 0 なら 0 をビットシフトで抽出
            // (符号ビットを取り出すだけで分岐命令は発生しない)
            let is_neg = (v >> (<$i>::BITS - 1)) as usize & 1;
            unsafe {
                // 先頭に '-' を書き込む。
                *ptr = b'-';
                // buf.add(is_neg) により、
                // - 正の数なら buf そのままの位置から書き始める (先頭の '-' は上書きされて消える)
                // - 負の数なら buf+1 の位置から書き始める (先頭の '-' が残る)
                is_neg + $g(ptr.add(is_neg), v.unsigned_abs())
            }
        }
    };
}
impl_ifmt!(i32, u32, di32, du32);
impl_ifmt!(i64, u64, di64, du64);
impl_ifmt!(i128, u128, di128, du128);

// ============================================================================
// FastFormat トレイトの実装 (u128 までの全整数型に対して)
// ============================================================================
/// Fast integer formatter trait used by [`FBufWriter`] integer APIs.
///
/// Primitive integer types up to `u128`/`i128` already implement this trait.
/// In most cases, you use it indirectly via methods such as [`FBufWriter::i`],
/// [`FBufWriter::spi`], [`FBufWriter::iln`], and [`FBufWriter::iv`].
///
/// # Examples
///
/// Basic usage with built-in integer types:
///
/// ```rust
/// use fastout::FBufWriter;
///
/// let mut out = FBufWriter::new(Vec::new());
/// out.i(42u32)?;
/// out.spi(-7i64)?;
/// out.ln()?;
///
/// let bytes = out.into_inner();
/// assert_eq!(String::from_utf8(bytes).unwrap(), "42 -7\n");
/// # std::io::Result::Ok(())
/// ```
///
/// Implementing `FastFormat` for your own wrapper type:
///
/// ```rust
/// use fastout::{FastFormat, FBufWriter};
///
/// struct Score(u32);
///
/// impl FastFormat for Score {
///     unsafe fn fast_format(&self, ptr: *mut u8) -> usize {
///         // Delegate formatting to the inner primitive integer.
///         self.0.fast_format(ptr)
///     }
/// }
///
/// let mut out = FBufWriter::new(Vec::new());
/// out.iln(Score(9001))?;
/// assert_eq!(String::from_utf8(out.into_inner()).unwrap(), "9001\n");
/// # std::io::Result::Ok(())
/// ```
pub trait FastFormat {
    /// Writes decimal ASCII representation into `ptr` and returns bytes written.
    ///
    /// # Safety
    ///
    /// `ptr` must point to writable memory with enough space for the largest
    /// possible output (`i128::MIN`, up to 40 bytes including `-`).
    unsafe fn fast_format(&self, ptr: *mut u8) -> usize;
}
macro_rules! impl_ffmt {
    ($t:ty, $i:ty, $f:ident) => {
        impl FastFormat for $t {
            unsafe fn fast_format(&self, ptr: *mut u8) -> usize {
                unsafe { $f(ptr, *self as $i) }
            }
        }
    };
}
impl_ffmt!(u8, u32, du32);
impl_ffmt!(u16, u32, du32);
impl_ffmt!(u32, u32, du32);
impl_ffmt!(u64, u64, du64);
impl_ffmt!(u128, u128, du128);
impl_ffmt!(usize, u64, du64);
impl_ffmt!(i8, i32, di32);
impl_ffmt!(i16, i32, di32);
impl_ffmt!(i32, i32, di32);
impl_ffmt!(i64, i64, di64);
impl_ffmt!(i128, i128, di128);
impl_ffmt!(isize, i64, di64);

// Buffer size for the custom writer
const SBUF: usize = 8192;
// Reserve space for i128 and delimiter
const SLIM: usize = SBUF - 48;
/// Buffered writer optimized for fast integer and byte/string output.
///
/// `FBufWriter` keeps an internal fixed-size buffer and minimizes `Write` calls.
/// Integer methods (`i`, `spi`, `iln`, `iv`) use [`FastFormat`] to avoid
/// `format!` allocations.
///
/// # Examples
///
/// ```rust
/// use fastout::FBufWriter;
///
/// let mut out = FBufWriter::new(Vec::new());
///
/// out.i(10u32)?;
/// out.spi(20u32)?;
/// out.spi(-30i64)?;
/// out.ln()?;
///
/// out.iv(&[1u32, 2, 3])?;
/// out.yes_ln(true)?;
/// out.str("done")?;
/// out.ln()?;
///
/// let bytes = out.into_inner();
/// assert_eq!(String::from_utf8(bytes).unwrap(), "10 20 -30\n1 2 3\nYes\ndone\n");
/// # std::io::Result::Ok(())
/// ```
pub struct FBufWriter<W: Write> {
    inner: Option<W>,
    buf: Box<[mem::MaybeUninit<u8>; SBUF]>,
    limit: *mut u8,
    p: *mut u8,
}
impl<W: Write> Write for FBufWriter<W> {
    fn write(&mut self, b: &[u8]) -> io::Result<usize> {
        let bl = b.len();
        let space = unsafe { self.limit.offset_from(self.p) as usize };
        if bl <= SBUF - SLIM {
            // If the input is small, write it directly to the buffer
            unsafe {
                c(b.as_ptr(), self.p, bl);
                self.p = self.p.add(bl);
            }
            if self.p > self.limit {
                self.flush()?
            }
        } else if SLIM < bl {
            // If the input is larger than the limit, flush the buffer first
            self.flush()?;
            // Write the remaining data directly to the inner writer
            self.inner.as_mut().unwrap().write_all(b)?
        } else if space < bl {
            // If there is not enough space, split the input
            unsafe {
                c(b.as_ptr(), self.p, space);
                self.p = self.limit;
            }
            self.flush()?;
            let b = &b[space..];
            // Write the remaining part of the input
            unsafe {
                c(b.as_ptr(), self.p, b.len());
                self.p = self.p.add(b.len())
            }
        } else {
            // If there is enough space, write to the buffer
            unsafe {
                c(b.as_ptr(), self.p, bl);
                self.p = self.p.add(bl)
            }
        }
        Ok(bl)
    }
    fn flush(&mut self) -> io::Result<()> {
        let base = self.buf.as_mut_ptr() as *mut u8;
        if self.p > base {
            unsafe {
                let len = self.p.offset_from(base) as usize;
                self.inner
                    .as_mut()
                    .unwrap()
                    .write_all(slice::from_raw_parts(base, len))?;
                self.p = base; // Reset pointer for the next chunk
            }
        }
        Ok(())
    }
}
impl<W: Write> Drop for FBufWriter<W> {
    fn drop(&mut self) {
        // Ensure the buffer is flushed on drop
        let _ = self.flush();
    }
}
impl<W: Write> FBufWriter<W> {
    pub fn new(inner: W) -> Self {
        let mut buf = Box::new([mem::MaybeUninit::uninit(); SBUF]);
        unsafe {
            let base = buf.as_mut_ptr() as *mut u8;
            FBufWriter {
                inner: Some(inner),
                buf,
                limit: base.add(SLIM),
                p: base,
            }
        }
    }
    pub fn into_inner(mut self) -> W {
        // Ensure the buffer is flushed before consuming the writer
        self.flush().unwrap();
        // Consume the writer and return the inner writer
        self.inner.take().unwrap()
    }
    #[inline]
    pub fn i<T: FastFormat>(&mut self, x: T) -> R {
        unsafe {
            self.p = self.p.add(x.fast_format(self.p));
        }
        if self.p > self.limit {
            self.flush()?
        }
        Ok(())
    }
    #[inline]
    pub fn spi<T: FastFormat>(&mut self, x: T) -> R {
        let mut p = self.p;
        unsafe {
            *p = b' ';
            p = p.add(1);
            p = p.add(x.fast_format(p));
        }
        self.p = p;
        if p > self.limit {
            self.flush()?
        }
        Ok(())
    }
    #[inline]
    pub fn isp<T: FastFormat>(&mut self, x: T) -> R {
        let mut p = self.p;
        unsafe {
            p = p.add(x.fast_format(p));
            *p = b' ';
            p = p.add(1);
        }
        self.p = p;
        if p > self.limit {
            self.flush()?
        }
        Ok(())
    }
    #[inline]
    pub fn iln<T: FastFormat>(&mut self, x: T) -> R {
        let mut p = self.p;
        unsafe {
            p = p.add(x.fast_format(p));
            *p = b'\n';
            p = p.add(1);
        }
        self.p = p;
        if p > self.limit {
            self.flush()?
        }
        Ok(())
    }
    pub fn iv<T: FastFormat + Copy>(&mut self, x: &[T]) -> R {
        if !x.is_empty() {
            self.i(x[0])?;
            for &y in x[1..].iter() {
                self.spi(y)?;
            }
        }
        unsafe {
            *self.p = b'\n';
            self.p = self.p.add(1);
        }
        if self.p > self.limit {
            self.flush()?
        }
        Ok(())
    }
    #[inline]
    pub fn byte(&mut self, b: u8) -> R {
        unsafe {
            *self.p = b;
            self.p = self.p.add(1);
        }
        if self.p > self.limit {
            self.flush()?;
        }
        Ok(())
    }
    /// Writes a short byte slice to the buffer, flushing if necessary
    #[inline]
    fn sbytes(&mut self, b: &[u8]) -> R {
        unsafe {
            assert_unchecked_pseudo(b.len() <= SBUF - SLIM);
            c(b.as_ptr(), self.p, b.len());
            self.p = self.p.add(b.len());
        }
        if self.p > self.limit {
            self.flush()?;
        }
        Ok(())
    }
    /// Writes a byte slice to the buffer, flushing if necessary
    #[inline]
    pub fn bytes(&mut self, b: &[u8]) -> R {
        // The write method processes all input bytes at once (unless an error occurs)
        self.write(b).map(|_| ())
    }
    /// Write the byte slice and then write a newline character
    pub fn bytes_ln(&mut self, b: &[u8]) -> R {
        self.bytes(b)?;
        self.ln()
    }
    /// Writes a Displayable object to the buffery
    pub fn fmt<T: fmt::Display>(&mut self, o: &T) -> R {
        self.bytes(o.to_string().as_bytes())
    }
    /// Writes a Displayable object followed by a newline to the buffer
    pub fn fmt_ln<T: fmt::Display>(&mut self, o: &T) -> R {
        self.bytes_ln(o.to_string().as_bytes())
    }
    /// Writes a str to the buffer, flushing if necessary
    #[inline]
    pub fn str(&mut self, s: &str) -> R {
        // Convert the string to bytes and write it to the buffer
        self.bytes(s.as_bytes())
    }
    /// Writes a space character to the buffer
    #[inline]
    pub fn sp(&mut self) -> R {
        self.byte(b' ')
    }
    /// Writes a newline character to the buffer
    #[inline]
    pub fn ln(&mut self) -> R {
        self.byte(b'\n')
    }
    /// Writes "Yes" or "No" to the buffer based on a boolean
    #[inline]
    pub fn yes(&mut self, x: bool) -> R {
        if x {
            self.sbytes(b"Yes")
        } else {
            self.sbytes(b"No")
        }
    }
    /// Writes "Yes" or "No" followed by a newline to the buffer based on a boolean
    #[inline]
    pub fn yes_ln(&mut self, x: bool) -> R {
        if x {
            self.sbytes(b"Yes\n")
        } else {
            self.sbytes(b"No\n")
        }
    }
}