Skip to main content

ferray_strings/
concat.rs

1// ferray-strings: Concatenation and repetition operations (REQ-3, REQ-4)
2//
3// Implements add (elementwise concat with broadcasting) and multiply (repeat).
4
5use ferray_core::dimension::{Dimension, IxDyn};
6use ferray_core::error::FerrayResult;
7
8use crate::string_array::{StringArray, broadcast_binary};
9
10/// Elementwise string concatenation with broadcasting.
11///
12/// Concatenates corresponding elements of `a` and `b`. If shapes differ,
13/// NumPy-style broadcasting is applied (e.g., a scalar string is broadcast
14/// against an array).
15///
16/// The result is always a dynamic-rank `StringArray<IxDyn>`.
17///
18/// # Errors
19/// Returns `FerrayError::BroadcastFailure` if shapes are incompatible.
20pub fn add<Da: Dimension, Db: Dimension>(
21    a: &StringArray<Da>,
22    b: &StringArray<Db>,
23) -> FerrayResult<StringArray<IxDyn>> {
24    let (out_shape, pairs) = broadcast_binary(a, b)?;
25    let a_data = a.as_slice();
26    let b_data = b.as_slice();
27
28    let data: Vec<String> = pairs
29        .iter()
30        .map(|&(ia, ib)| format!("{}{}", a_data[ia], b_data[ib]))
31        .collect();
32
33    StringArray::from_vec(IxDyn::new(&out_shape), data)
34}
35
36/// Same-dimension elementwise string concatenation.
37///
38/// Like [`add`] but both inputs must have the same shape — no
39/// broadcasting is performed, and the result preserves the static
40/// dimension type. Use this when you know the shapes match and want
41/// to keep `StringArray<Ix1>` instead of getting `StringArray<IxDyn>`
42/// (#163).
43///
44/// # Errors
45/// Returns `FerrayError::ShapeMismatch` if shapes differ.
46pub fn add_same<D: Dimension>(
47    a: &StringArray<D>,
48    b: &StringArray<D>,
49) -> FerrayResult<StringArray<D>> {
50    if a.shape() != b.shape() {
51        return Err(ferray_core::error::FerrayError::shape_mismatch(format!(
52            "add_same: shapes {:?} and {:?} must be identical",
53            a.shape(),
54            b.shape()
55        )));
56    }
57    let data: Vec<String> = a
58        .iter()
59        .zip(b.iter())
60        .map(|(x, y)| format!("{x}{y}"))
61        .collect();
62    StringArray::from_vec(a.dim().clone(), data)
63}
64
65/// Repeat each string element `n` times.
66///
67/// # Errors
68/// Returns an error if the internal array construction fails.
69pub fn multiply<D: Dimension>(a: &StringArray<D>, n: usize) -> FerrayResult<StringArray<D>> {
70    a.map(|s| s.repeat(n))
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use crate::string_array::array;
77
78    #[test]
79    fn test_add_same_shape() {
80        let a = array(&["hello", "foo"]).unwrap();
81        let b = array(&[" world", " bar"]).unwrap();
82        let c = add(&a, &b).unwrap();
83        assert_eq!(c.as_slice(), &["hello world", "foo bar"]);
84    }
85
86    #[test]
87    fn test_add_broadcast_scalar() {
88        // AC-2: strings::add broadcasts a scalar string against an array correctly
89        let a = array(&["hello", "world"]).unwrap();
90        let b = array(&["!"]).unwrap();
91        let c = add(&a, &b).unwrap();
92        assert_eq!(c.as_slice(), &["hello!", "world!"]);
93    }
94
95    #[test]
96    fn test_add_broadcast_scalar_left() {
97        let a = array(&[">> "]).unwrap();
98        let b = array(&["hello", "world"]).unwrap();
99        let c = add(&a, &b).unwrap();
100        assert_eq!(c.as_slice(), &[">> hello", ">> world"]);
101    }
102
103    #[test]
104    fn test_add_incompatible_shapes() {
105        let a = array(&["a", "b", "c"]).unwrap();
106        let b = array(&["x", "y"]).unwrap();
107        assert!(add(&a, &b).is_err());
108    }
109
110    #[test]
111    fn test_multiply() {
112        let a = array(&["ab", "cd"]).unwrap();
113        let b = multiply(&a, 3).unwrap();
114        assert_eq!(b.as_slice(), &["ababab", "cdcdcd"]);
115    }
116
117    #[test]
118    fn test_multiply_zero() {
119        let a = array(&["hello"]).unwrap();
120        let b = multiply(&a, 0).unwrap();
121        assert_eq!(b.as_slice(), &[""]);
122    }
123
124    #[test]
125    fn test_multiply_one() {
126        let a = array(&["hello"]).unwrap();
127        let b = multiply(&a, 1).unwrap();
128        assert_eq!(b.as_slice(), &["hello"]);
129    }
130}