use std::collections::HashMap;
use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::{OnceLock, RwLock};
use unicode_width::{UnicodeWidthChar, UnicodeWidthStr};
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum WidthMethod {
#[default]
WcWidth,
Unicode,
}
const WIDTH_METHOD_WCWIDTH: u8 = 0;
const WIDTH_METHOD_UNICODE: u8 = 1;
static WIDTH_METHOD: AtomicU8 = AtomicU8::new(WIDTH_METHOD_WCWIDTH);
static WIDTH_OVERRIDES: OnceLock<RwLock<HashMap<char, usize>>> = OnceLock::new();
static WIDTH_OVERRIDES_ENABLED: std::sync::atomic::AtomicBool =
std::sync::atomic::AtomicBool::new(false);
fn width_overrides() -> &'static RwLock<HashMap<char, usize>> {
WIDTH_OVERRIDES.get_or_init(|| RwLock::new(HashMap::new()))
}
pub fn set_width_override(ch: char, width: usize) {
{
let mut map = width_overrides()
.write()
.expect("width override lock poisoned");
map.insert(ch, width);
}
WIDTH_OVERRIDES_ENABLED.store(true, Ordering::Release);
}
#[must_use]
pub fn get_width_override(ch: char) -> Option<usize> {
if !WIDTH_OVERRIDES_ENABLED.load(Ordering::Acquire) {
return None;
}
let map = WIDTH_OVERRIDES
.get()?
.read()
.expect("width override lock poisoned");
map.get(&ch).copied()
}
pub fn clear_width_overrides() {
if let Some(map) = WIDTH_OVERRIDES.get() {
map.write().expect("width override lock poisoned").clear();
}
WIDTH_OVERRIDES_ENABLED.store(false, Ordering::Release);
}
pub fn set_width_method(method: WidthMethod) {
let value = match method {
WidthMethod::WcWidth => WIDTH_METHOD_WCWIDTH,
WidthMethod::Unicode => WIDTH_METHOD_UNICODE,
};
WIDTH_METHOD.store(value, Ordering::Relaxed);
}
#[must_use]
pub fn width_method() -> WidthMethod {
match WIDTH_METHOD.load(Ordering::Relaxed) {
WIDTH_METHOD_UNICODE => WidthMethod::Unicode,
_ => WidthMethod::WcWidth,
}
}
#[must_use]
pub fn display_width(s: &str) -> usize {
display_width_with_method(s, width_method())
}
#[inline]
#[must_use]
pub fn display_width_char(c: char) -> usize {
if c.is_ascii() && (' '..='~').contains(&c) {
return 1;
}
if c < ' ' {
return 0;
}
display_width_char_with_method(c, width_method())
}
#[must_use]
pub fn display_width_with_method(s: &str, method: WidthMethod) -> usize {
if WIDTH_OVERRIDES_ENABLED.load(Ordering::Acquire) {
return s
.chars()
.map(|ch| display_width_char_with_method(ch, method))
.sum();
}
match method {
WidthMethod::WcWidth => UnicodeWidthStr::width(s),
WidthMethod::Unicode => UnicodeWidthStr::width_cjk(s),
}
}
#[must_use]
pub fn display_width_char_with_method(c: char, method: WidthMethod) -> usize {
if let Some(width) = get_width_override(c) {
return width;
}
match method {
WidthMethod::WcWidth => UnicodeWidthChar::width(c).unwrap_or(0),
WidthMethod::Unicode => UnicodeWidthChar::width_cjk(c).unwrap_or(0),
}
}
#[must_use]
pub fn is_zero_width(c: char) -> bool {
display_width_char(c) == 0
}
#[must_use]
pub fn is_wide(c: char) -> bool {
display_width_char(c) == 2
}
#[cfg(test)]
mod tests {
use super::*;
struct ClearOverridesOnDrop;
impl Drop for ClearOverridesOnDrop {
fn drop(&mut self) {
clear_width_overrides();
}
}
#[test]
fn test_ascii_width() {
assert_eq!(display_width("hello"), 5);
assert_eq!(display_width_char('a'), 1);
}
#[test]
fn test_cjk_width() {
assert_eq!(display_width("æ¼¢å—"), 4);
assert_eq!(display_width_char('æ¼¢'), 2);
assert!(is_wide('æ¼¢'));
}
#[test]
fn test_emoji_width() {
assert_eq!(display_width("😀"), 2);
}
#[test]
fn test_zero_width() {
assert!(is_zero_width('\u{0301}')); }
#[test]
fn test_width_methods() {
let ch = 'â‘ ';
assert_eq!(display_width_char_with_method(ch, WidthMethod::WcWidth), 1);
assert_eq!(display_width_char_with_method(ch, WidthMethod::Unicode), 2);
}
#[test]
fn test_width_overrides_set_get_clear() {
let _guard = ClearOverridesOnDrop;
assert_eq!(get_width_override('🦀'), None);
set_width_override('🦀', 1);
assert_eq!(get_width_override('🦀'), Some(1));
clear_width_overrides();
assert_eq!(get_width_override('🦀'), None);
}
#[test]
fn test_width_calculation_uses_override() {
let _guard = ClearOverridesOnDrop;
assert_eq!(display_width_char('🦀'), 2);
set_width_override('🦀', 1);
assert_eq!(display_width_char('🦀'), 1);
assert_eq!(display_width("A🦀B"), 3);
}
}