<commit_before>use crate::{
    decimal::Decimal,
    duration::Duration,
    errors::{AvroResult, Error},
    schema::Schema,
    types::Value,
    util::{safe_len, zag_i32, zag_i64},
};
use std::{collections::HashMap, io::Read, str::FromStr};
use uuid::Uuid;

#[inline]
fn decode_long<R: Read>(reader: &mut R) -> AvroResult<Value> {
    zag_i64(reader).map(Value::Long)
}

#[inline]
fn decode_int<R: Read>(reader: &mut R) -> AvroResult<Value> {
    zag_i32(reader).map(Value::Int)
}

#[inline]
fn decode_len<R: Read>(reader: &mut R) -> AvroResult<usize> {
    zag_i64(reader).and_then(|len| safe_len(len as usize))
}

/// Decode a `Value` from avro format given its `Schema`.
pub fn decode<R: Read>(schema: &Schema, reader: &mut R) -> AvroResult<Value> {
    match *schema {
        Schema::Null => Ok(Value::Null),
        Schema::Boolean => {
            let mut buf = [0u8; 1];
            reader.read_exact(&mut buf[..])?;

            match buf[0] {
                0u8 => Ok(Value::Boolean(false)),
                1u8 => Ok(Value::Boolean(true)),
                _ => Err(Error::Decode("not a bool".to_string())),
            }
        }
        Schema::Decimal { ref inner, .. } => match **inner {
            Schema::Fixed { .. } => match decode(inner, reader)? {
                Value::Fixed(_, bytes) => Ok(Value::Decimal(Decimal::from(bytes))),
                _ => Err(Error::Decode(
                    "not a fixed value, required for decimal with fixed schema".to_string(),
                )),
            },
            Schema::Bytes => match decode(inner, reader)? {
                Value::Bytes(bytes) => Ok(Value::Decimal(Decimal::from(bytes))),
                _ => Err(Error::Decode(
                    "not a bytes value, required for decimal with bytes schema".to_string(),
                )),
            },
            _ => Err(Error::Decode(
                "not a fixed or bytes type, required for decimal schema".to_string(),
            )),
        },
        Schema::Uuid => Ok(Value::Uuid(Uuid::from_str(
            match decode(&Schema::String, reader)? {
                Value::String(ref s) => s,
                _ => {
                    return Err(Error::Decode(
                        "not a string type, required for uuid".to_string(),
                    ))
                }
            },
        )?)),
        Schema::Int => decode_int(reader),
        Schema::Date => zag_i32(reader).map(Value::Date),
        Schema::TimeMillis => zag_i32(reader).map(Value::TimeMillis),
        Schema::Long => decode_long(reader),
        Schema::TimeMicros => zag_i64(reader).map(Value::TimeMicros),
        Schema::TimestampMillis => zag_i64(reader).map(Value::TimestampMillis),
        Schema::TimestampMicros => zag_i64(reader).map(Value::TimestampMicros),
        Schema::Duration => {
            let mut buf = [0u8; 12];
            reader.read_exact(&mut buf)?;
            Ok(Value::Duration(Duration::from(buf)))
        }
        Schema::Float => {
            let mut buf = [0u8; std::mem::size_of::<f32>()];
            reader.read_exact(&mut buf[..])?;
            Ok(Value::Float(f32::from_le_bytes(buf)))
        }
        Schema::Double => {
            let mut buf = [0u8; std::mem::size_of::<f64>()];
            reader.read_exact(&mut buf[..])?;
            Ok(Value::Double(f64::from_le_bytes(buf)))
        }
        Schema::Bytes => {
            let len = decode_len(reader)?;
            let mut buf = vec![0u8; len];
            reader.read_exact(&mut buf)?;
            Ok(Value::Bytes(buf))
        }
        Schema::String => {
            let len = decode_len(reader)?;
            let mut buf = vec![0u8; len];
            reader.read_exact(&mut buf)?;

            String::from_utf8(buf)
                .map(Value::String)
                .map_err(|_| Error::Decode("not a valid utf-8 string".to_string()))
        }
        Schema::Fixed { size, .. } => {
            let mut buf = vec![0u8; size as usize];
            reader.read_exact(&mut buf)?;
            Ok(Value::Fixed(size, buf))
        }
        Schema::Array(ref inner) => {
            let mut items = Vec::new();

            loop {
                let raw_len = zag_i64(reader)?;
                // arrays are 0-terminated, 0i64 is also encoded as 0 in Avro
                // reading a length of 0 means the end of the array
                let len = safe_len(match raw_len.cmp(&0) {
                    std::cmp::Ordering::Equal => break,
                    std::cmp::Ordering::Less => {
                        let _size = zag_i64(reader)?;
                        -raw_len
                    }
                    std::cmp::Ordering::Greater => raw_len,
                } as usize)?;

                items.reserve(len as usize);
                for _ in 0..len {
                    items.push(decode(inner, reader)?);
                }
            }

            Ok(Value::Array(items))
        }
        Schema::Map(ref inner) => {
            let mut items = HashMap::new();

            loop {
                let raw_len = zag_i64(reader)?;
                // maps are 0-terminated, 0i64 is also encoded as 0 in Avro
                // reading a length of 0 means the end of the map
                let len = safe_len(match raw_len.cmp(&0) {
                    std::cmp::Ordering::Equal => break,
                    std::cmp::Ordering::Less => {
                        let _size = zag_i64(reader)?;
                        -raw_len
                    }
                    std::cmp::Ordering::Greater => raw_len,
                } as usize)?;

                items.reserve(len);
                for _ in 0..len {
                    if let Value::String(key) = decode(&Schema::String, reader)? {
                        let value = decode(inner, reader)?;
                        items.insert(key, value);
                    } else {
                        return Err(Error::Decode("map key is not a string".to_string()));
                    }
                }
            }

            Ok(Value::Map(items))
        }
        Schema::Union(ref inner) => {
            let index = zag_i64(reader)?;
            let variants = inner.variants();
            let variant = variants
                .get(index as usize)
                .ok_or_else(|| Error::Decode("Union index out of bounds".to_string()))?;
            let value = decode(variant, reader)?;
            Ok(Value::Union(Box::new(value)))
        }
        Schema::Record { ref fields, .. } => {
            // Benchmarks indicate ~10% improvement using this method.
            let mut items = Vec::with_capacity(fields.len());
            for field in fields {
                // This clone is also expensive. See if we can do away with it...
                items.push((field.name.clone(), decode(&field.schema, reader)?));
            }
            Ok(Value::Record(items))
        }
        Schema::Enum { ref symbols, .. } => {
            if let Value::Int(index) = decode_int(reader)? {
                if index >= 0 && (index as usize) <= symbols.len() {
                    let symbol = symbols[index as usize].clone();
                    Ok(Value::Enum(index, symbol))
                } else {
                    Err(Error::Decode("enum symbol index out of bounds".to_string()))
                }
            } else {
                Err(Error::Decode("enum symbol not found".to_string()))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::Value::{Array, Int, Map};

    #[test]
    fn test_decode_array_without_size() {
        let mut input: &[u8] = &[6, 2, 4, 6, 0];
        let result = decode(&Schema::Array(Box::new(Schema::Int)), &mut input);
        assert_eq!(Array(vec!(Int(1), Int(2), Int(3))), result.unwrap());
    }

    #[test]
    fn test_decode_array_with_size() {
        let mut input: &[u8] = &[5, 6, 2, 4, 6, 0];
        let result = decode(&Schema::Array(Box::new(Schema::Int)), &mut input);
        assert_eq!(Array(vec!(Int(1), Int(2), Int(3))), result.unwrap());
    }

    #[test]
    fn test_decode_map_without_size() {
        let mut input: &[u8] = &[0x02, 0x08, 0x74, 0x65, 0x73, 0x74, 0x02, 0x00];
        let result = decode(&Schema::Map(Box::new(Schema::Int)), &mut input);
        let mut expected = HashMap::new();
        expected.insert(String::from("test"), Int(1));
        assert_eq!(Map(expected), result.unwrap());
    }

    #[test]
    fn test_decode_map_with_size() {
        let mut input: &[u8] = &[0x01, 0x0C, 0x08, 0x74, 0x65, 0x73, 0x74, 0x02, 0x00];
        let result = decode(&Schema::Map(Box::new(Schema::Int)), &mut input);
        let mut expected = HashMap::new();
        expected.insert(String::from("test"), Int(1));
        assert_eq!(Map(expected), result.unwrap());
    }

    #[test]
    fn test_negative_decimal_value() {
        use crate::{encode::encode, schema::Name};
        use num_bigint::ToBigInt;
        let inner = Box::new(Schema::Fixed {
            size: 2,
            name: Name::new("decimal"),
        });
        let schema = Schema::Decimal {
            inner,
            precision: 4,
            scale: 2,
        };
        let bigint = -423.to_bigint().unwrap();
        let value = Value::Decimal(Decimal::from(bigint.to_signed_bytes_be()));

        let mut buffer = Vec::new();
        encode(&value, &schema, &mut buffer);

        let mut bytes = &buffer[..];
        let result = decode(&schema, &mut bytes).unwrap();
        assert_eq!(result, value);
    }

    #[test]
    fn test_decode_decimal_with_bigger_than_necessary_size() {
        use crate::{encode::encode, schema::Name};
        use num_bigint::ToBigInt;
        let inner = Box::new(Schema::Fixed {
            size: 13,
            name: Name::new("decimal"),
        });
        let schema = Schema::Decimal {
            inner,
            precision: 4,
            scale: 2,
        };
        let value = Value::Decimal(Decimal::from(
            (-423.to_bigint().unwrap()).to_signed_bytes_be(),
        ));
        let mut buffer = Vec::<u8>::new();

        encode(&value, &schema, &mut buffer);
        let mut bytes: &[u8] = &buffer[..];
        let result = decode(&schema, &mut bytes).unwrap();
        assert_eq!(result, value);
    }
}
<commit_msg>Fix length casting during decoding (#145)<commit_after>use crate::{
    decimal::Decimal,
    duration::Duration,
    errors::{AvroResult, Error},
    schema::Schema,
    types::Value,
    util::{safe_len, zag_i32, zag_i64},
};
use std::{collections::HashMap, convert::TryFrom, io::Read, str::FromStr};
use uuid::Uuid;

#[inline]
fn decode_long<R: Read>(reader: &mut R) -> AvroResult<Value> {
    zag_i64(reader).map(Value::Long)
}

#[inline]
fn decode_int<R: Read>(reader: &mut R) -> AvroResult<Value> {
    zag_i32(reader).map(Value::Int)
}

#[inline]
fn decode_len<R: Read>(reader: &mut R) -> AvroResult<usize> {
    safe_len(usize::try_from(zag_i64(reader)?)?)
}

/// Decode the length of a sequence.
///
/// Maps and arrays are 0-terminated, 0i64 is also encoded as 0 in Avro reading a length of 0 means
/// the end of the map or array.
fn decode_seq_len<R: Read>(reader: &mut R) -> AvroResult<usize> {
    let raw_len = zag_i64(reader)?;
    safe_len(usize::try_from(match raw_len.cmp(&0) {
        std::cmp::Ordering::Equal => return Ok(0),
        std::cmp::Ordering::Less => {
            let _size = zag_i64(reader)?;
            -raw_len
        }
        std::cmp::Ordering::Greater => raw_len,
    })?)
}

/// Decode a `Value` from avro format given its `Schema`.
pub fn decode<R: Read>(schema: &Schema, reader: &mut R) -> AvroResult<Value> {
    match *schema {
        Schema::Null => Ok(Value::Null),
        Schema::Boolean => {
            let mut buf = [0u8; 1];
            reader.read_exact(&mut buf[..])?;

            match buf[0] {
                0u8 => Ok(Value::Boolean(false)),
                1u8 => Ok(Value::Boolean(true)),
                _ => Err(Error::Decode("not a bool".to_string())),
            }
        }
        Schema::Decimal { ref inner, .. } => match **inner {
            Schema::Fixed { .. } => match decode(inner, reader)? {
                Value::Fixed(_, bytes) => Ok(Value::Decimal(Decimal::from(bytes))),
                _ => Err(Error::Decode(
                    "not a fixed value, required for decimal with fixed schema".to_string(),
                )),
            },
            Schema::Bytes => match decode(inner, reader)? {
                Value::Bytes(bytes) => Ok(Value::Decimal(Decimal::from(bytes))),
                _ => Err(Error::Decode(
                    "not a bytes value, required for decimal with bytes schema".to_string(),
                )),
            },
            _ => Err(Error::Decode(
                "not a fixed or bytes type, required for decimal schema".to_string(),
            )),
        },
        Schema::Uuid => Ok(Value::Uuid(Uuid::from_str(
            match decode(&Schema::String, reader)? {
                Value::String(ref s) => s,
                _ => {
                    return Err(Error::Decode(
                        "not a string type, required for uuid".to_string(),
                    ))
                }
            },
        )?)),
        Schema::Int => decode_int(reader),
        Schema::Date => zag_i32(reader).map(Value::Date),
        Schema::TimeMillis => zag_i32(reader).map(Value::TimeMillis),
        Schema::Long => decode_long(reader),
        Schema::TimeMicros => zag_i64(reader).map(Value::TimeMicros),
        Schema::TimestampMillis => zag_i64(reader).map(Value::TimestampMillis),
        Schema::TimestampMicros => zag_i64(reader).map(Value::TimestampMicros),
        Schema::Duration => {
            let mut buf = [0u8; 12];
            reader.read_exact(&mut buf)?;
            Ok(Value::Duration(Duration::from(buf)))
        }
        Schema::Float => {
            let mut buf = [0u8; std::mem::size_of::<f32>()];
            reader.read_exact(&mut buf[..])?;
            Ok(Value::Float(f32::from_le_bytes(buf)))
        }
        Schema::Double => {
            let mut buf = [0u8; std::mem::size_of::<f64>()];
            reader.read_exact(&mut buf[..])?;
            Ok(Value::Double(f64::from_le_bytes(buf)))
        }
        Schema::Bytes => {
            let len = decode_len(reader)?;
            let mut buf = vec![0u8; len];
            reader.read_exact(&mut buf)?;
            Ok(Value::Bytes(buf))
        }
        Schema::String => {
            let len = decode_len(reader)?;
            let mut buf = vec![0u8; len];
            reader.read_exact(&mut buf)?;

            String::from_utf8(buf)
                .map(Value::String)
                .map_err(|_| Error::Decode("not a valid utf-8 string".to_string()))
        }
        Schema::Fixed { size, .. } => {
            let mut buf = vec![0u8; size];
            reader.read_exact(&mut buf)?;
            Ok(Value::Fixed(size, buf))
        }
        Schema::Array(ref inner) => {
            let mut items = Vec::new();

            loop {
                let len = decode_seq_len(reader)?;
                if len == 0 {
                    break;
                }

                items.reserve(len);
                for _ in 0..len {
                    items.push(decode(inner, reader)?);
                }
            }

            Ok(Value::Array(items))
        }
        Schema::Map(ref inner) => {
            let mut items = HashMap::new();

            loop {
                let len = decode_seq_len(reader)?;
                if len == 0 {
                    break;
                }

                items.reserve(len);
                for _ in 0..len {
                    if let Value::String(key) = decode(&Schema::String, reader)? {
                        let value = decode(inner, reader)?;
                        items.insert(key, value);
                    } else {
                        return Err(Error::Decode("map key is not a string".to_string()));
                    }
                }
            }

            Ok(Value::Map(items))
        }
        Schema::Union(ref inner) => {
            let index = zag_i64(reader)?;
            let variants = inner.variants();
            let variant = variants
                .get(usize::try_from(index)?)
                .ok_or_else(|| Error::Decode("Union index out of bounds".to_string()))?;
            let value = decode(variant, reader)?;
            Ok(Value::Union(Box::new(value)))
        }
        Schema::Record { ref fields, .. } => {
            // Benchmarks indicate ~10% improvement using this method.
            let mut items = Vec::with_capacity(fields.len());
            for field in fields {
                // TODO: This clone is also expensive. See if we can do away with it...
                items.push((field.name.clone(), decode(&field.schema, reader)?));
            }
            Ok(Value::Record(items))
        }
        Schema::Enum { ref symbols, .. } => {
            if let Value::Int(raw_index) = decode_int(reader)? {
                let index = usize::try_from(raw_index)?;
                if (0..=symbols.len()).contains(&index) {
                    let symbol = symbols[index].clone();
                    Ok(Value::Enum(raw_index, symbol))
                } else {
                    Err(Error::Decode("enum symbol index out of bounds".to_string()))
                }
            } else {
                Err(Error::Decode("enum symbol not found".to_string()))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{decode, Decimal, Schema, Value};
    use crate::types::Value::{Array, Int, Map};
    use std::collections::HashMap;

    #[test]
    fn test_decode_array_without_size() {
        let mut input: &[u8] = &[6, 2, 4, 6, 0];
        let result = decode(&Schema::Array(Box::new(Schema::Int)), &mut input);
        assert_eq!(Array(vec!(Int(1), Int(2), Int(3))), result.unwrap());
    }

    #[test]
    fn test_decode_array_with_size() {
        let mut input: &[u8] = &[5, 6, 2, 4, 6, 0];
        let result = decode(&Schema::Array(Box::new(Schema::Int)), &mut input);
        assert_eq!(Array(vec!(Int(1), Int(2), Int(3))), result.unwrap());
    }

    #[test]
    fn test_decode_map_without_size() {
        let mut input: &[u8] = &[0x02, 0x08, 0x74, 0x65, 0x73, 0x74, 0x02, 0x00];
        let result = decode(&Schema::Map(Box::new(Schema::Int)), &mut input);
        let mut expected = HashMap::new();
        expected.insert(String::from("test"), Int(1));
        assert_eq!(Map(expected), result.unwrap());
    }

    #[test]
    fn test_decode_map_with_size() {
        let mut input: &[u8] = &[0x01, 0x0C, 0x08, 0x74, 0x65, 0x73, 0x74, 0x02, 0x00];
        let result = decode(&Schema::Map(Box::new(Schema::Int)), &mut input);
        let mut expected = HashMap::new();
        expected.insert(String::from("test"), Int(1));
        assert_eq!(Map(expected), result.unwrap());
    }

    #[test]
    fn test_negative_decimal_value() {
        use crate::{encode::encode, schema::Name};
        use num_bigint::ToBigInt;
        let inner = Box::new(Schema::Fixed {
            size: 2,
            name: Name::new("decimal"),
        });
        let schema = Schema::Decimal {
            inner,
            precision: 4,
            scale: 2,
        };
        let bigint = -423.to_bigint().unwrap();
        let value = Value::Decimal(Decimal::from(bigint.to_signed_bytes_be()));

        let mut buffer = Vec::new();
        encode(&value, &schema, &mut buffer);

        let mut bytes = &buffer[..];
        let result = decode(&schema, &mut bytes).unwrap();
        assert_eq!(result, value);
    }

    #[test]
    fn test_decode_decimal_with_bigger_than_necessary_size() {
        use crate::{encode::encode, schema::Name};
        use num_bigint::ToBigInt;
        let inner = Box::new(Schema::Fixed {
            size: 13,
            name: Name::new("decimal"),
        });
        let schema = Schema::Decimal {
            inner,
            precision: 4,
            scale: 2,
        };
        let value = Value::Decimal(Decimal::from(
            (-423.to_bigint().unwrap()).to_signed_bytes_be(),
        ));
        let mut buffer = Vec::<u8>::new();

        encode(&value, &schema, &mut buffer);
        let mut bytes: &[u8] = &buffer[..];
        let result = decode(&schema, &mut bytes).unwrap();
        assert_eq!(result, value);
    }
}
<|endoftext|>


<commit_before>//! Decodes a Bitterlemon-encoded byte stream into its original bit stream.

use std::iter::Iterator;

pub fn decode<S>(input: S) -> Decoder<S>
where S : Iterator<Item=u8> {
	Decoder::<S> {
		source: input,
		state: DecodeState::Pending,
	}
}

pub struct Decoder<S> {
	state: DecodeState,
	source: S,
}

#[derive(Debug, PartialEq, Eq)]
pub enum Error {
	TruncatedInput(u8, u8),
}

type IterationResult = Result<bool, Error>;

impl<S> Iterator for Decoder<S>
where S : Iterator<Item=u8> {
	type Item = IterationResult;

	fn next(&mut self) -> Option<Self::Item> {

		// pull from source if needs be
		if self.must_pull() {
			let next = self.source.next();
			self.next_with_pulled(next)
		} else {
			self.next_from_existing()
		}
	}
}

impl<S> Decoder<S>
where S : Iterator<Item=u8> {

	fn must_pull(&self) -> bool {
		match self.state {
			DecodeState::Pending => true,
			DecodeState::Done => false,
			DecodeState::Run(remaining, _) => remaining == 0,
			DecodeState::Frame(remaining, _, stage_size) => remaining == 0 || stage_size == 0,
		}
	}

	fn next_with_pulled(&mut self, next: Option<u8>) -> Option<<Self as Iterator>::Item> {

		// handle None from source
		let next = match next {
			Some(x) => x,
			None => match self.state {
				DecodeState::Pending => {
					self.state = DecodeState::Done;
					return None;
				}, // source was empty
				DecodeState::Done    => { return None; }, // always return None here
				DecodeState::Run(_, _) => {
					unreachable!("next_with_pulled called with more run bits to flush: {:?}", self.state);
				},
				DecodeState::Frame(remaining, _, stage_size) => {
					debug_assert!(stage_size == 0);
					debug_assert!(remaining > 0);

					// missing bytes to complete the frame
					let error_specifics = Error::TruncatedInput(remaining, (remaining + 7) >> 3);
					return Some(Err(error_specifics));
				}
			}
		};

		// handle mid-frame
		if match self.state {
			DecodeState::Frame(ref mut remaining, ref mut stage, ref mut stage_size)
			if *remaining > 0 => {
				debug_assert!(*stage_size == 0); // shouldn't have pulled otherwise
				*stage = next;
				*stage_size = 8;
				// now fall through to real iteration logic
				true
			},
			_ => false
		} {
			return self.next_from_existing();
		}

		let got = match next {
			n if n < 0x80 => {
				// frame beginning
				let frame_size = byte_to_frame_size(n);
				self.state = DecodeState::Frame(frame_size, 0, 0);
				None
			},
			n => {
				// new run
				let frame_size = byte_to_run_size(n);
				let mode = n >= 0xc0;
				self.state = if frame_size > 1 {
					// don't bother moving to run state if only one pixel in this run
					// also, leaving this method in state Run(0, _) is a logic error
					DecodeState::Run(frame_size - 1, mode)
				}
				else {
					DecodeState::Pending
				};
				Some(Ok(mode))
			}
		};

		got.or_else(|| {
			let next = self.source.next();
			self.next_with_pulled(next)
		})
	}

	fn next_from_existing(&mut self) -> Option<<Self as Iterator>::Item> {

		let (to_return, next_state) = match self.state {
			DecodeState::Pending => unreachable!(),
			DecodeState::Done    => { return None; },
			DecodeState::Run(ref mut remaining, ref run_mode) => {
				*remaining -= 1;
				(*run_mode, if *remaining == 0 {Some(DecodeState::Pending)} else {None})
			},
			DecodeState::Frame(ref mut remaining, ref mut stage, ref mut stage_size) => {
				let got_bit = (*stage & 0x80) != 0;
				*stage = (*stage & 0x7f) << 1;
				*stage_size -= 1;
				*remaining -= 1;
				(got_bit, if *remaining == 0 {Some(DecodeState::Pending)} else {None})
			}
		};

		if let Some(next_state) = next_state {
			self.state = next_state;
		}

		Some(Ok(to_return))
	}

}

fn byte_to_run_size(byte: u8) -> u8 {
	let byte = byte & 0x3f;
	if byte == 0 { 0x40 } else { byte }
}

fn byte_to_frame_size(byte: u8) -> u8 {
	if byte == 0 { 0x80 } else { byte }
}

#[derive(Debug)]
enum DecodeState {
	Pending, // should pull
	Run(u8, bool), // run count, is_set
	Frame(u8, u8, u8), // frame pixel count, stage contents, stage size
	Done, // iteration complete
}

#[cfg(test)]
mod test_decoder {

	macro_rules! decoder {
		( $($e:expr),* ) => {
			{
				let v = vec![$( $e, )*];
				super::decode(v.into_iter())
			}
		}
	}

	#[test]
	fn empty_input() {
		let mut iter = decoder![];

		assert_eq!(iter.next(), None);
		assert_eq!(iter.next(), None);
	}

	fn single_run_impl(top_bits: u8, mode: bool) {
		for i in 0..0x3fu8 {
			let run_size = super::byte_to_run_size(i);
			let mut iter = decoder![i+top_bits];
			for _ in 0..run_size {
				assert_eq!(iter.next(), Some(Ok(mode)));
			}
			assert_eq!(iter.next(), None);
		}
	}

	#[test]
	fn single_run_clear() {
		single_run_impl(0x80, false)
	}

	#[test]
	fn single_run_set() {
		single_run_impl(0xc0, true)
	}

	#[test]
	fn single_byte_frame() {
		let case = |byte_in: u8, bool_out: bool| {
			let mut iter = decoder![0x01, byte_in];
			assert_eq!(iter.next(), Some(Ok(bool_out)));
		};

		case(0xff, true);
		case(0x00, false);
		case(0x80, true);
		case(0x7f, false);
	}

	#[test]
	fn full_byte_frame() {
		let mut iter = decoder![0x08, 0x55];

		let mut expected = false;
		for _ in 0..8 {
			assert_eq!(iter.next(), Some(Ok(expected)));
			expected = !expected;
		}

		assert_eq!(iter.next(), None);
	}

	#[test]
	fn two_byte_frame() {
		let mut iter = decoder![0x0f, 0x55, 0x55];

		let mut expected = false;
		for _ in 0..15 {
			assert_eq!(iter.next(), Some(Ok(expected)));
			expected = !expected;
		}

		assert_eq!(iter.next(), None);
	}

	#[test]
	fn alternate_runs_frames() {
		let case = |bytes: &[u8], count: usize, first_output: bool| {
			let mut iter = super::decode(bytes.iter().map(|&b| b));

			let mut expected = first_output;
			for _ in 0..count {
				assert_eq!(iter.next(), Some(Ok(expected)));
				expected = !expected;
			}

			assert_eq!(iter.next(), None);
		};

		case(&[0xc1, 0x10, 0x55, 0x55, 0x81], 18, true);
		case(&[0x81, 0x10, 0xaa, 0xaa, 0xc1], 18, false);
		case(&[0x08, 0xaa, 0xc1, 0x08, 0x55], 17, true);
		case(&[0x08, 0x55, 0x81, 0x08, 0xaa], 17, false);
	}

	#[test]
	fn error_on_frame_cutoff() {
		let case = |bytes: &[u8], pixels_lost: u8, bytes_missing: u8| {
			let mut iter = super::decode(bytes.iter().map(|&b| b));

			let ok_count = (bytes.len() - 1) * 8;
			for _ in 0..ok_count {
				assert!(match iter.next() {
					Some(Ok(_)) => true,
					_ => false
				});
			}

			let error = iter.next();
			assert!(error.is_some());
			let error = error.unwrap();

			assert!(error.is_err());
			match error.unwrap_err() {
				super::Error::TruncatedInput(pl, bm) => {
					assert_eq!(pl, pixels_lost);
					assert_eq!(bm, bytes_missing);
				}
			};
		};

		case(&[0x01], 1, 1);
		case(&[0x02], 2, 1);
		case(&[0x00], 0x80, 0x10);
		case(&[0x09, 0xff], 1, 1);
		case(&[0x7f, 0,1,2,3,4,5,6,7,8,9,10,11,12,13,14], 7, 1);
	}
}
<commit_msg>Decoder idempotence test<commit_after>//! Decodes a Bitterlemon-encoded byte stream into its original bit stream.

use std::iter::Iterator;

pub fn decode<S>(input: S) -> Decoder<S>
where S : Iterator<Item=u8> {
	Decoder::<S> {
		source: input,
		state: DecodeState::Pending,
	}
}

pub struct Decoder<S> {
	state: DecodeState,
	source: S,
}

#[derive(Debug, PartialEq, Eq)]
pub enum Error {
	TruncatedInput(u8, u8),
}

type IterationResult = Result<bool, Error>;

impl<S> Iterator for Decoder<S>
where S : Iterator<Item=u8> {
	type Item = IterationResult;

	fn next(&mut self) -> Option<Self::Item> {

		// pull from source if needs be
		if self.must_pull() {
			let next = self.source.next();
			self.next_with_pulled(next)
		} else {
			self.next_from_existing()
		}
	}
}

impl<S> Decoder<S>
where S : Iterator<Item=u8> {

	fn must_pull(&self) -> bool {
		match self.state {
			DecodeState::Pending => true,
			DecodeState::Done => false,
			DecodeState::Run(remaining, _) => remaining == 0,
			DecodeState::Frame(remaining, _, stage_size) => remaining == 0 || stage_size == 0,
		}
	}

	fn next_with_pulled(&mut self, next: Option<u8>) -> Option<<Self as Iterator>::Item> {

		// handle None from source
		let next = match next {
			Some(x) => x,
			None => match self.state {
				DecodeState::Pending => {
					self.state = DecodeState::Done;
					return None;
				}, // source was empty
				DecodeState::Done    => { return None; }, // always return None here
				DecodeState::Run(_, _) => {
					unreachable!("next_with_pulled called with more run bits to flush: {:?}", self.state);
				},
				DecodeState::Frame(remaining, _, stage_size) => {
					debug_assert!(stage_size == 0);
					debug_assert!(remaining > 0);

					// missing bytes to complete the frame
					let error_specifics = Error::TruncatedInput(remaining, (remaining + 7) >> 3);
					return Some(Err(error_specifics));
				}
			}
		};

		// handle mid-frame
		if match self.state {
			DecodeState::Frame(ref mut remaining, ref mut stage, ref mut stage_size)
			if *remaining > 0 => {
				debug_assert!(*stage_size == 0); // shouldn't have pulled otherwise
				*stage = next;
				*stage_size = 8;
				// now fall through to real iteration logic
				true
			},
			_ => false
		} {
			return self.next_from_existing();
		}

		let got = match next {
			n if n < 0x80 => {
				// frame beginning
				let frame_size = byte_to_frame_size(n);
				self.state = DecodeState::Frame(frame_size, 0, 0);
				None
			},
			n => {
				// new run
				let frame_size = byte_to_run_size(n);
				let mode = n >= 0xc0;
				self.state = if frame_size > 1 {
					// don't bother moving to run state if only one pixel in this run
					// also, leaving this method in state Run(0, _) is a logic error
					DecodeState::Run(frame_size - 1, mode)
				}
				else {
					DecodeState::Pending
				};
				Some(Ok(mode))
			}
		};

		got.or_else(|| {
			let next = self.source.next();
			self.next_with_pulled(next)
		})
	}

	fn next_from_existing(&mut self) -> Option<<Self as Iterator>::Item> {

		let (to_return, next_state) = match self.state {
			DecodeState::Pending => unreachable!(),
			DecodeState::Done    => { return None; },
			DecodeState::Run(ref mut remaining, ref run_mode) => {
				*remaining -= 1;
				(*run_mode, if *remaining == 0 {Some(DecodeState::Pending)} else {None})
			},
			DecodeState::Frame(ref mut remaining, ref mut stage, ref mut stage_size) => {
				let got_bit = (*stage & 0x80) != 0;
				*stage = (*stage & 0x7f) << 1;
				*stage_size -= 1;
				*remaining -= 1;
				(got_bit, if *remaining == 0 {Some(DecodeState::Pending)} else {None})
			}
		};

		if let Some(next_state) = next_state {
			self.state = next_state;
		}

		Some(Ok(to_return))
	}

}

fn byte_to_run_size(byte: u8) -> u8 {
	let byte = byte & 0x3f;
	if byte == 0 { 0x40 } else { byte }
}

fn byte_to_frame_size(byte: u8) -> u8 {
	if byte == 0 { 0x80 } else { byte }
}

#[derive(Debug)]
enum DecodeState {
	Pending, // should pull
	Run(u8, bool), // run count, is_set
	Frame(u8, u8, u8), // frame pixel count, stage contents, stage size
	Done, // iteration complete
}

#[cfg(test)]
mod test_decoder {

	macro_rules! decoder {
		( $($e:expr),* ) => {
			{
				let v = vec![$( $e, )*];
				super::decode(v.into_iter())
			}
		}
	}

	#[test]
	fn empty_input() {
		let mut iter = decoder![];

		assert_eq!(iter.next(), None);
		assert_eq!(iter.next(), None);
	}

	fn single_run_impl(top_bits: u8, mode: bool) {
		for i in 0..0x3fu8 {
			let run_size = super::byte_to_run_size(i);
			let mut iter = decoder![i+top_bits];
			for _ in 0..run_size {
				assert_eq!(iter.next(), Some(Ok(mode)));
			}
			assert_eq!(iter.next(), None);
		}
	}

	#[test]
	fn single_run_clear() {
		single_run_impl(0x80, false)
	}

	#[test]
	fn single_run_set() {
		single_run_impl(0xc0, true)
	}

	#[test]
	fn single_byte_frame() {
		let case = |byte_in: u8, bool_out: bool| {
			let mut iter = decoder![0x01, byte_in];
			assert_eq!(iter.next(), Some(Ok(bool_out)));
		};

		case(0xff, true);
		case(0x00, false);
		case(0x80, true);
		case(0x7f, false);
	}

	#[test]
	fn full_byte_frame() {
		let mut iter = decoder![0x08, 0x55];

		let mut expected = false;
		for _ in 0..8 {
			assert_eq!(iter.next(), Some(Ok(expected)));
			expected = !expected;
		}

		assert_eq!(iter.next(), None);
	}

	#[test]
	fn two_byte_frame() {
		let mut iter = decoder![0x0f, 0x55, 0x55];

		let mut expected = false;
		for _ in 0..15 {
			assert_eq!(iter.next(), Some(Ok(expected)));
			expected = !expected;
		}

		assert_eq!(iter.next(), None);
	}

	#[test]
	fn alternate_runs_frames() {
		let case = |bytes: &[u8], count: usize, first_output: bool| {
			let mut iter = super::decode(bytes.iter().map(|&b| b));

			let mut expected = first_output;
			for _ in 0..count {
				assert_eq!(iter.next(), Some(Ok(expected)));
				expected = !expected;
			}

			assert_eq!(iter.next(), None);
		};

		case(&[0xc1, 0x10, 0x55, 0x55, 0x81], 18, true);
		case(&[0x81, 0x10, 0xaa, 0xaa, 0xc1], 18, false);
		case(&[0x08, 0xaa, 0xc1, 0x08, 0x55], 17, true);
		case(&[0x08, 0x55, 0x81, 0x08, 0xaa], 17, false);
	}

	#[test]
	fn error_on_frame_cutoff() {
		let case = |bytes: &[u8], pixels_lost: u8, bytes_missing: u8| {
			let mut iter = super::decode(bytes.iter().map(|&b| b));

			let ok_count = (bytes.len() - 1) * 8;
			for _ in 0..ok_count {
				assert!(match iter.next() {
					Some(Ok(_)) => true,
					_ => false
				});
			}

			let error = iter.next();
			assert!(error.is_some());
			let error = error.unwrap();

			assert!(error.is_err());
			match error.unwrap_err() {
				super::Error::TruncatedInput(pl, bm) => {
					assert_eq!(pl, pixels_lost);
					assert_eq!(bm, bytes_missing);
				}
			};
		};

		case(&[0x01], 1, 1);
		case(&[0x02], 2, 1);
		case(&[0x00], 0x80, 0x10);
		case(&[0x09, 0xff], 1, 1);
		case(&[0x7f, 0,1,2,3,4,5,6,7,8,9,10,11,12,13,14], 7, 1);
	}

	#[test]
	fn next_none_idempotence() {
		let src = &[0xc1u8];
		let mut iter = super::decode(src.iter().map(|&b| b));

		assert!(iter.next().is_some());

		for _ in 0..20 {
			assert_eq!(iter.next(), None);
		}
	}
}
<|endoftext|>


<commit_before>//! Utilities for creating zero-cost wakers that don't do anything.
use futures_core::task::{LocalWaker, UnsafeWake, Waker};
use core::ptr::NonNull;

#[derive(Debug)]
struct NoopWake {
    _reserved: (),
}

unsafe impl UnsafeWake for NoopWake {
    unsafe fn clone_raw(&self) -> Waker {
        noop_waker()
    }

    unsafe fn drop_raw(&self) {}

    unsafe fn wake(&self) {}
}

fn noop_unsafe_wake() -> NonNull<dyn UnsafeWake> {
    static mut INSTANCE: NoopWake = NoopWake { _reserved: () };
    unsafe { NonNull::new_unchecked(&mut INSTANCE as *mut dyn UnsafeWake) }
}

fn noop_waker() -> Waker {
    unsafe { Waker::new(noop_unsafe_wake()) }
}

/// Create a new [`LocalWaker`](futures_core::task::LocalWaker) referencing a
/// singleton instance of [`NoopWake`].
#[inline]
pub fn noop_local_waker() -> LocalWaker {
    unsafe { LocalWaker::new(noop_unsafe_wake()) }
}

/// Get a thread local reference to a
/// [`LocalWaker`](futures_core::task::LocalWaker) referencing a singleton
/// instance of [`NoopWake`].
///
/// # Examples
///
/// ```
/// #![feature(futures_api)]
/// use futures::task::noop_local_waker_ref;
/// let lw = noop_local_waker_ref();
/// lw.wake();
/// ```
#[inline]
pub fn noop_local_waker_ref() -> &'static LocalWaker {
    static NOOP_WAKE_REF: &(dyn UnsafeWake + Sync) = &NoopWake { _reserved: () };
    // Unsafety: `Waker` and `LocalWaker` are `repr(transparent)` wrappers around
    // `NonNull<dyn UnsafeWake>`, which has the same repr as `&(dyn UnsafeWake + Sync)`
    // So an &'static &(dyn UnsafeWake + Sync) can be unsafely cast to a
    // &'static LocalWaker
    unsafe {
        core::mem::transmute::<
            &&(dyn UnsafeWake + Sync),
            &'static LocalWaker,
        >(&NOOP_WAKE_REF)
    }
}
<commit_msg>select!: calm clippy<commit_after>//! Utilities for creating zero-cost wakers that don't do anything.
use futures_core::task::{LocalWaker, UnsafeWake, Waker};
use core::ptr::NonNull;

#[derive(Debug)]
struct NoopWake {
    _reserved: (),
}

unsafe impl UnsafeWake for NoopWake {
    unsafe fn clone_raw(&self) -> Waker {
        noop_waker()
    }

    unsafe fn drop_raw(&self) {}

    unsafe fn wake(&self) {}
}

fn noop_unsafe_wake() -> NonNull<dyn UnsafeWake> {
    static mut INSTANCE: NoopWake = NoopWake { _reserved: () };
    unsafe { NonNull::new_unchecked(&mut INSTANCE as *mut dyn UnsafeWake) }
}

fn noop_waker() -> Waker {
    unsafe { Waker::new(noop_unsafe_wake()) }
}

/// Create a new [`LocalWaker`](futures_core::task::LocalWaker) referencing a
/// singleton instance of [`NoopWake`].
#[inline]
pub fn noop_local_waker() -> LocalWaker {
    unsafe { LocalWaker::new(noop_unsafe_wake()) }
}

/// Get a thread local reference to a
/// [`LocalWaker`](futures_core::task::LocalWaker) referencing a singleton
/// instance of [`NoopWake`].
///
/// # Examples
///
/// ```
/// #![feature(futures_api)]
/// use futures::task::noop_local_waker_ref;
/// let lw = noop_local_waker_ref();
/// lw.wake();
/// ```
#[inline]
pub fn noop_local_waker_ref() -> &'static LocalWaker {
    static NOOP_WAKE_REF: &(dyn UnsafeWake + Sync) = &NoopWake { _reserved: () };
    // Unsafety: `Waker` and `LocalWaker` are `repr(transparent)` wrappers around
    // `NonNull<dyn UnsafeWake>`, which has the same repr as `&(dyn UnsafeWake + Sync)`
    // So an &'static &(dyn UnsafeWake + Sync) can be unsafely cast to a
    // &'static LocalWaker
    #[allow(clippy::transmute_ptr_to_ptr)]
    unsafe {
        core::mem::transmute::<
            &&(dyn UnsafeWake + Sync),
            &'static LocalWaker,
        >(&NOOP_WAKE_REF)
    }
}
<|endoftext|>


<commit_before>use cocoa::base::{id, nil};
use cocoa::foundation::{NSUInteger, NSString};
use sys::{MTLCreateSystemDefaultDevice, MTLDevice};
use std::borrow::Cow;
use std::error::Error;
use std::fmt::{self, Display, Formatter};
use std::ffi::CStr;

pub struct Device(id);

impl Device {
    pub fn system_default_device() -> Result<Self, DeviceError> {
        let device = unsafe { MTLCreateSystemDefaultDevice() };
        if device != nil {
            Ok(Device(device))
        } else {
            Err(DeviceError::ConstructionFailed)
        }
    }

    pub unsafe fn get_raw(&self) -> id {
        self.0
    }

    pub fn is_headless(&self) -> bool {
        unsafe { self.0.headless() != 0 }
    }

    pub fn is_low_power(&self) -> bool {
        unsafe { self.0.lowPower() != 0 }
    }

    pub fn name(&self) -> Cow<str> {
        unsafe { CStr::from_ptr(self.0.name().UTF8String()).to_string_lossy() }
    }

    pub fn supports_texture_sample_count(&self, sample_count: usize) -> bool {
        unsafe { self.0.supportsTextureSampleCount(sample_count as NSUInteger) != 0 }
    }
}

/// Internal utility function to create a Device without exposing internals.
/// Not exported publicly from this crate.
#[doc(hidden)]
pub fn _make_device(device: id) -> Device {
    Device(device)
}

#[derive(Clone, Debug)]
pub enum DeviceError {
    ConstructionFailed
}

impl Display for DeviceError {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let descr = match *self {
            DeviceError::ConstructionFailed => "DeviceError::ConstructionFailed"
        };
        write!(f, "{}", descr)
    }
}

impl Error for DeviceError {
    fn description(&self) -> &str {
        match *self {
            DeviceError::ConstructionFailed =>
                "Could not create a default device. Please ensure that you are using at least OSX 10.11 or iOS 8.0"
        }
    }
}
<commit_msg>Implement `Device::max_threads_per_group`<commit_after>use cocoa::base::{id, nil};
use cocoa::foundation::{NSUInteger, NSString};
use sys::{MTLCreateSystemDefaultDevice, MTLDevice};
use std::borrow::Cow;
use std::error::Error;
use std::fmt::{self, Display, Formatter};
use std::ffi::CStr;
use {Size};

pub struct Device(id);

impl Device {
    pub fn system_default_device() -> Result<Self, DeviceError> {
        let device = unsafe { MTLCreateSystemDefaultDevice() };
        if device != nil {
            Ok(Device(device))
        } else {
            Err(DeviceError::ConstructionFailed)
        }
    }

    pub unsafe fn get_raw(&self) -> id {
        self.0
    }

    pub fn is_headless(&self) -> bool {
        unsafe { self.0.headless() != 0 }
    }

    pub fn is_low_power(&self) -> bool {
        unsafe { self.0.lowPower() != 0 }
    }

    pub fn max_threads_per_group(&self) -> Size {
        unsafe { self.0.maxThreadsPerGroup() }
    }

    pub fn name(&self) -> Cow<str> {
        unsafe { CStr::from_ptr(self.0.name().UTF8String()).to_string_lossy() }
    }

    pub fn supports_texture_sample_count(&self, sample_count: usize) -> bool {
        unsafe { self.0.supportsTextureSampleCount(sample_count as NSUInteger) != 0 }
    }
}

/// Internal utility function to create a Device without exposing internals.
/// Not exported publicly from this crate.
#[doc(hidden)]
pub fn _make_device(device: id) -> Device {
    Device(device)
}

#[derive(Clone, Debug)]
pub enum DeviceError {
    ConstructionFailed
}

impl Display for DeviceError {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let descr = match *self {
            DeviceError::ConstructionFailed => "DeviceError::ConstructionFailed"
        };
        write!(f, "{}", descr)
    }
}

impl Error for DeviceError {
    fn description(&self) -> &str {
        match *self {
            DeviceError::ConstructionFailed =>
                "Could not create a default device. Please ensure that you are using at least OSX 10.11 or iOS 8.0"
        }
    }
}
<|endoftext|>


<commit_before>use constants::{CR, DEFAULT_LINE_SIZE, ESCAPE, LF, NUL};
use crc32;

use std::fs::File;
use std::io;
use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write};
use std::path::Path;

/// Options for encoding
#[derive(Debug)]
pub struct EncodeOptions {
    line_length: u8,
    parts: u32,
    part: u32,
    begin: u64,
    end: u64,
}

impl Default for EncodeOptions {
    /// Constructs a new EncodeOptions instance, with the following defaults:
    /// line_length = 128.
    /// parts = 1,
    /// part = begin = end = 0
    fn default() -> Self {
        EncodeOptions {
            line_length: DEFAULT_LINE_SIZE,
            parts: 1,
            part: 0,
            begin: 0,
            end: 0,
        }
    }
}

impl EncodeOptions {
    pub fn new() -> EncodeOptions {
        Default::default()
    }

    /// Sets the maximum line length.
    pub fn line_length(mut self, line_length: u8) -> EncodeOptions {
        self.line_length = line_length;
        self
    }

    /// Sets the number of parts (default=1).
    /// When the number of parts equals 1, no '=ypart' line will be written
    /// in the ouput.
    pub fn parts(mut self, parts: u32) -> EncodeOptions {
        self.parts = parts;
        self
    }

    /// Sets the part number.
    /// Only used when `parts > 1`.
    /// The part number count starts at 1.
    pub fn part(mut self, part: u32) -> EncodeOptions {
        self.part = part;
        self
    }

    /// Sets the begin (which is the file offset + 1).
    /// Only used when `parts > 1`.
    /// The size of the part is `end - begin + 1`.
    pub fn begin(mut self, begin: u64) -> EncodeOptions {
        self.begin = begin;
        self
    }

    /// Sets the end.
    /// Only used when `parts > 1`.
    /// The size of the part is `end - begin + 1`.
    /// `end` should be larger than `begin`, otherwise an overflow error occurrs.
    pub fn end(mut self, end: u64) -> EncodeOptions {
        self.end = end;
        self
    }
}

/// Encodes the input file in a new output file.
/// # Example
/// ```rust,no_run
/// let encode_options = yenc::EncodeOptions::default().parts(1);
/// let mut output_file = std::fs::File::create("test1.bin.yenc").unwrap();
/// yenc::encode_file("test1.bin", &encode_options, &mut output_file);
/// ```
/// # Errors
/// - when the output file already exists
///
pub fn encode_file<P, W>(
    input_path: P,
    encode_options: &EncodeOptions,
    output: W,
) -> Result<(), io::Error>
where
    P: AsRef<Path>,
    W: Write,
{
    let input_filename = input_path.as_ref().file_name();
    let input_filename = match input_filename {
        Some(s) => s.to_str().unwrap_or(""),
        None => "",
    };
    let input_file = File::open(&input_path)?;
    let length = input_file.metadata()?.len();

    encode_stream(input_file, output, length, input_filename, encode_options)
}

pub fn encode_stream<R, W>(
    input: R,
    output: W,
    length: u64,
    input_filename: &str,
    encode_options: &EncodeOptions,
) -> io::Result<()>
where
    R: Read + Seek,
    W: Write,
{
    let mut rdr = BufReader::new(input);
    let mut checksum = crc32::Crc32::new();
    let mut buffer = [0u8; 8192];
    let mut col = 0;
    let mut output = BufWriter::new(output);

    if encode_options.parts == 1 {
        write!(
            output,
            "=ybegin line={} size={} name={}\r\n",
            encode_options.line_length, length, input_filename
        )?;
    } else {
        write!(
            output,
            "=ybegin part={} line={} size={} name={}\r\n",
            encode_options.part, encode_options.line_length, length, input_filename
        )?;
    }

    if encode_options.parts > 1 {
        write!(
            output,
            "=ypart begin={} end={}\r\n",
            encode_options.begin, encode_options.end
        )?;
    }

    rdr.seek(SeekFrom::Start(encode_options.begin - 1))?;

    let mut remainder = (encode_options.end - encode_options.begin + 1) as usize;
    while remainder > 0 {
        let buf_slice = if remainder > buffer.len() {
            &mut buffer[..]
        } else {
            &mut buffer[0..remainder]
        };
        rdr.read_exact(buf_slice)?;
        checksum.update_with_slice(buf_slice);
        match encode_buffer(buf_slice, col, encode_options.line_length, &mut output) {
            Ok(c) => col = c,
            Err(e) => return Err(e),
        };
        remainder -= buf_slice.len();
    }

    if encode_options.parts > 1 {
        write!(
            output,
            "\r\n=yend size={} part={} pcrc32={:08x}\r\n",
            checksum.num_bytes, encode_options.part, checksum.crc
        )?;
    } else {
        write!(
            output,
            "\r\n=yend size={} crc32={:08x}\r\n",
            checksum.num_bytes, checksum.crc
        )?;
    }
    Ok(())
}

/// Encode the byte slice into a vector of yEncoded bytes.
///
/// Lines are wrapped with a maximum of `line_length` characters per line.
/// Does not include the header and footer lines. These are only produced
/// by `encode_stream` and `encode_file`.
pub fn encode_buffer<W>(input: &[u8], col: u8, line_length: u8, writer: W) -> io::Result<u8>
where
    W: Write,
{
    let mut col = col;
    let mut writer = writer;
    let mut v = Vec::<u8>::with_capacity(input.len() * 102 / 100);
    input.iter().for_each(|&b| {
        let (encoded, encoded_len) = encode_byte(b);
        v.push(encoded.0);
        if encoded_len > 1 {
            v.push(encoded.1);
        }
        col += encoded_len as u8;
        if col >= line_length {
            v.push(CR);
            v.push(LF);
            col = 0;
        }
    });
    writer.write_all(&v)?;
    Ok(col)
}

#[inline(always)]
fn encode_byte(input_byte: u8) -> ((u8, u8), usize) {
    let mut output = (0, 0);

    let output_byte = input_byte.overflowing_add(42).0;
    let len = match output_byte {
        NUL | CR | LF | ESCAPE => {
            output.0 = ESCAPE;
            output.1 = output_byte.overflowing_add(64).0;
            2
        }
        _ => {
            output.0 = output_byte;
            1
        }
    };
    (output, len)
}

#[cfg(test)]
mod tests {
    use super::super::constants::{CR, ESCAPE, LF, NUL};
    use super::{encode_buffer, encode_byte};

    #[test]
    fn escape_null() {
        assert_eq!(((ESCAPE, 0x40), 2), encode_byte(214));
    }

    /*
    #[test]
    fn escape_tab() {
        let mut output = [0u8; 2];
        assert_eq!(2, encode_byte(214 + TAB, &mut output));
        assert_eq!(vec![ESCAPE, 0x49], output);
    }
    */

    #[test]
    fn escape_lf() {
        assert_eq!(((ESCAPE, 0x4A), 2), encode_byte(214 + LF));
    }

    #[test]
    fn escape_cr() {
        assert_eq!(((ESCAPE, 0x4D), 2), encode_byte(214 + CR));
    }

    /*
    #[test]
    fn escape_space() {
        let mut output = [0u8; 2];
        assert_eq!(2, encode_byte(214 + SPACE, &mut output));
        assert_eq!(vec![ESCAPE, 0x60], output);
    }
    */

    #[test]
    fn escape_equal_sign() {
        assert_eq!(((ESCAPE, 0x7D), 2), encode_byte(ESCAPE - 42));
    }

    #[test]
    fn non_escaped() {
        for x in 0..256u16 {
            let encoded = (x as u8).overflowing_add(42).0;
            if encoded != NUL && encoded != CR && encoded != LF && encoded != ESCAPE {
                assert_eq!(((encoded, 0), 1), encode_byte(x as u8));
            }
        }
    }

    #[test]
    fn test_encode_buffer() {
        let buffer = (0..256u16).map(|c| c as u8).collect::<Vec<u8>>();
        #[cfg_attr(rustfmt, rustfmt_skip)]
        const EXPECTED: [u8; 264] =
                       [42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61,
                       125, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80,
                       81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100,
                       101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116,
                       117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132,
                       133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148,
                       149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164,
                       165, 166, 167, 168, 13, 10, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178,
                       179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194,
                       195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210,
                       211, 212, 213, 214, 215, 216,217, 218, 219, 220, 221, 222, 223, 224, 225, 226,
                       227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242,
                       243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 61, 64, 1, 2, 3,
                       4, 5, 6, 7, 8, 9, 61, 74, 11, 12, 61, 77, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
                       24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 13, 10, 38, 39, 40, 41];
        let mut encoded = Vec::<u8>::new();
        let result = encode_buffer(&buffer, 0, 128, &mut encoded);
        assert!(result.is_ok());
        assert_eq!(encoded.as_slice(), &EXPECTED[..]);
    }
}
<commit_msg>Fix asymmetric yencode/ydecode (leading dots)<commit_after>use constants::{CR, DEFAULT_LINE_SIZE, DOT, ESCAPE, LF, NUL};
use crc32;

use std::fs::File;
use std::io;
use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write};
use std::path::Path;

/// Options for encoding
#[derive(Debug)]
pub struct EncodeOptions {
    line_length: u8,
    parts: u32,
    part: u32,
    begin: u64,
    end: u64,
}

impl Default for EncodeOptions {
    /// Constructs a new EncodeOptions instance, with the following defaults:
    /// line_length = 128.
    /// parts = 1,
    /// part = begin = end = 0
    fn default() -> Self {
        EncodeOptions {
            line_length: DEFAULT_LINE_SIZE,
            parts: 1,
            part: 0,
            begin: 0,
            end: 0,
        }
    }
}

impl EncodeOptions {
    pub fn new() -> EncodeOptions {
        Default::default()
    }

    /// Sets the maximum line length.
    pub fn line_length(mut self, line_length: u8) -> EncodeOptions {
        self.line_length = line_length;
        self
    }

    /// Sets the number of parts (default=1).
    /// When the number of parts equals 1, no '=ypart' line will be written
    /// in the ouput.
    pub fn parts(mut self, parts: u32) -> EncodeOptions {
        self.parts = parts;
        self
    }

    /// Sets the part number.
    /// Only used when `parts > 1`.
    /// The part number count starts at 1.
    pub fn part(mut self, part: u32) -> EncodeOptions {
        self.part = part;
        self
    }

    /// Sets the begin (which is the file offset + 1).
    /// Only used when `parts > 1`.
    /// The size of the part is `end - begin + 1`.
    pub fn begin(mut self, begin: u64) -> EncodeOptions {
        self.begin = begin;
        self
    }

    /// Sets the end.
    /// Only used when `parts > 1`.
    /// The size of the part is `end - begin + 1`.
    /// `end` should be larger than `begin`, otherwise an overflow error occurrs.
    pub fn end(mut self, end: u64) -> EncodeOptions {
        self.end = end;
        self
    }
}

/// Encodes the input file in a new output file.
/// # Example
/// ```rust,no_run
/// let encode_options = yenc::EncodeOptions::default().parts(1);
/// let mut output_file = std::fs::File::create("test1.bin.yenc").unwrap();
/// yenc::encode_file("test1.bin", &encode_options, &mut output_file);
/// ```
/// # Errors
/// - when the output file already exists
///
pub fn encode_file<P, W>(
    input_path: P,
    encode_options: &EncodeOptions,
    output: W,
) -> Result<(), io::Error>
where
    P: AsRef<Path>,
    W: Write,
{
    let input_filename = input_path.as_ref().file_name();
    let input_filename = match input_filename {
        Some(s) => s.to_str().unwrap_or(""),
        None => "",
    };
    let input_file = File::open(&input_path)?;
    let length = input_file.metadata()?.len();

    encode_stream(input_file, output, length, input_filename, encode_options)
}

pub fn encode_stream<R, W>(
    input: R,
    output: W,
    length: u64,
    input_filename: &str,
    encode_options: &EncodeOptions,
) -> io::Result<()>
where
    R: Read + Seek,
    W: Write,
{
    let mut rdr = BufReader::new(input);
    let mut checksum = crc32::Crc32::new();
    let mut buffer = [0u8; 8192];
    let mut col = 0;
    let mut output = BufWriter::new(output);

    if encode_options.parts == 1 {
        write!(
            output,
            "=ybegin line={} size={} name={}\r\n",
            encode_options.line_length, length, input_filename
        )?;
    } else {
        write!(
            output,
            "=ybegin part={} line={} size={} name={}\r\n",
            encode_options.part, encode_options.line_length, length, input_filename
        )?;
    }

    if encode_options.parts > 1 {
        write!(
            output,
            "=ypart begin={} end={}\r\n",
            encode_options.begin, encode_options.end
        )?;
    }

    rdr.seek(SeekFrom::Start(encode_options.begin - 1))?;

    let mut remainder = (encode_options.end - encode_options.begin + 1) as usize;
    while remainder > 0 {
        let buf_slice = if remainder > buffer.len() {
            &mut buffer[..]
        } else {
            &mut buffer[0..remainder]
        };
        rdr.read_exact(buf_slice)?;
        checksum.update_with_slice(buf_slice);
        match encode_buffer(buf_slice, col, encode_options.line_length, &mut output) {
            Ok(c) => col = c,
            Err(e) => return Err(e),
        };
        remainder -= buf_slice.len();
    }

    if encode_options.parts > 1 {
        write!(
            output,
            "\r\n=yend size={} part={} pcrc32={:08x}\r\n",
            checksum.num_bytes, encode_options.part, checksum.crc
        )?;
    } else {
        write!(
            output,
            "\r\n=yend size={} crc32={:08x}\r\n",
            checksum.num_bytes, checksum.crc
        )?;
    }
    Ok(())
}

/// Encode the byte slice into a vector of yEncoded bytes.
///
/// Lines are wrapped with a maximum of `line_length` characters per line.
/// Does not include the header and footer lines. These are only produced
/// by `encode_stream` and `encode_file`.
pub fn encode_buffer<W>(input: &[u8], col: u8, line_length: u8, writer: W) -> io::Result<u8>
where
    W: Write,
{
    let mut col = col;
    let mut writer = writer;
    let mut v = Vec::<u8>::with_capacity(input.len() * 104 / 100);
    input.iter().for_each(|&b| {
        let encoded = encode_byte(b);
        v.push(encoded.0);
        col += match encoded.0 {
            ESCAPE => {
                v.push(encoded.1);
                2
            }
            DOT if col == 0 => {
                v.push(DOT);
                2
            }
            _ => 1,
        };
        if col >= line_length {
            v.push(CR);
            v.push(LF);
            col = 0;
        }
    });
    writer.write_all(&v)?;
    Ok(col)
}

#[inline(always)]
fn encode_byte(input_byte: u8) -> (u8, u8) {
    let mut output = (0, 0);

    let output_byte = input_byte.overflowing_add(42).0;
    match output_byte {
        NUL | CR | LF | ESCAPE => {
            output.0 = ESCAPE;
            output.1 = output_byte.overflowing_add(64).0;
        }
        _ => {
            output.0 = output_byte;
        }
    };
    output
}

#[cfg(test)]
mod tests {
    use super::super::constants::{CR, ESCAPE, LF, NUL};
    use super::{encode_buffer, encode_byte};

    #[test]
    fn escape_null() {
        assert_eq!((ESCAPE, 0x40), encode_byte(214));
    }

    /*
    #[test]
    fn escape_tab() {
        let mut output = [0u8; 2];
        assert_eq!(2, encode_byte(214 + TAB, &mut output));
        assert_eq!(vec![ESCAPE, 0x49], output);
    }
    */

    #[test]
    fn escape_lf() {
        assert_eq!((ESCAPE, 0x4A), encode_byte(214 + LF));
    }

    #[test]
    fn escape_cr() {
        assert_eq!((ESCAPE, 0x4D), encode_byte(214 + CR));
    }

    /*
    #[test]
    fn escape_space() {
        let mut output = [0u8; 2];
        assert_eq!(2, encode_byte(214 + SPACE, &mut output));
        assert_eq!(vec![ESCAPE, 0x60], output);
    }
    */

    #[test]
    fn escape_equal_sign() {
        assert_eq!((ESCAPE, 0x7D), encode_byte(ESCAPE - 42));
    }

    #[test]
    fn non_escaped() {
        for x in 0..256u16 {
            let encoded = (x as u8).overflowing_add(42).0;
            if encoded != NUL && encoded != CR && encoded != LF && encoded != ESCAPE {
                assert_eq!((encoded, 0), encode_byte(x as u8));
            }
        }
    }

    #[test]
    fn test_encode_buffer() {
        let buffer = (0..256u16).map(|c| c as u8).collect::<Vec<u8>>();
        #[cfg_attr(rustfmt, rustfmt_skip)]
        const EXPECTED: [u8; 264] =
                       [42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61,
                       125, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80,
                       81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100,
                       101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116,
                       117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132,
                       133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148,
                       149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164,
                       165, 166, 167, 168, 13, 10, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178,
                       179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194,
                       195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210,
                       211, 212, 213, 214, 215, 216,217, 218, 219, 220, 221, 222, 223, 224, 225, 226,
                       227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242,
                       243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 61, 64, 1, 2, 3,
                       4, 5, 6, 7, 8, 9, 61, 74, 11, 12, 61, 77, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
                       24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 13, 10, 38, 39, 40, 41];
        let mut encoded = Vec::<u8>::new();
        let result = encode_buffer(&buffer, 0, 128, &mut encoded);
        assert!(result.is_ok());
        assert_eq!(encoded.as_slice(), &EXPECTED[..]);
    }
}
<|endoftext|>


<commit_before>use std::cmp::min;
use std::iter::IteratorExt;
use std::collections::{BTreeMap, RingBuf};
use std::string::ToString;
use serialize::json::Json;

use self::TemplateElement::{RawString, Expression, HelperExpression,
                            HTMLExpression, HelperBlock, Comment};

#[deriving(PartialEq)]
pub struct Template {
    pub elements: Vec<TemplateElement>
}

#[deriving(PartialEq, Show)]
enum ParserState {
    Text,
    HtmlExpression,
    Comment,
    HelperStart,
    HelperEnd,
    Expression,
    Invalid
}

#[deriving(PartialEq)]
pub struct Helper {
    name: String,
    params: Vec<String>,
    hash: BTreeMap<String, Json>,
    template: Option<Template>,
    inverse: Option<Template>,
    block: bool
}

impl Helper {
    pub fn name(&self) -> &String {
        &self.name
    }

    pub fn params(&self) -> &Vec<String> {
        &self.params
    }

    pub fn hash(&self) -> &BTreeMap<String, Json> {
        &self.hash
    }

    pub fn template(&self) -> Option<&Template> {
        match self.template {
            Some(ref t) => {
                Some(t)
            },
            None => None
        }
    }

    pub fn inverse(&self) -> Option<&Template> {
        match self.inverse {
            Some(ref t) => {
                Some(t)
            },
            None => None
        }
    }

    pub fn is_block(&self) -> bool {
        self.block
    }
}

impl ToString for Helper {
    fn to_string(&self) -> String {

        let mut buf = String::new();

        if self.block {
            buf.push_str(format!("{{{{#{}", self.name).as_slice());
        } else {
            buf.push_str(format!("{{{{{}", self.name).as_slice());
        }

        for p in self.params.iter() {
            buf.push_str(format!(" {}", p).as_slice());
        }

        for k in self.hash.keys() {
            buf.push_str(format!(" {}={}", k, self.hash.get(k).unwrap()).as_slice());
        }

        buf.push_str("}}");

        if self.block {
            let tpl = self.template();
            if tpl.is_some() {
                buf.push_str(tpl.unwrap().to_string().as_slice());
            }
            let ivs = self.inverse();
            if ivs.is_some() {
                buf.push_str("{{else}}");
                buf.push_str(ivs.unwrap().to_string().as_slice());
            }
            buf.push_str(format!("{{{{/{}}}}}", self.name).as_slice());
        }
        buf
    }
}

#[deriving(Show, Copy)]
pub struct TemplateError;

impl Helper {
    fn parse(source: String, block: bool) -> Result<Helper, TemplateError> {
        let mut tokens = source.split(|&: c: char| c.is_whitespace())
            .filter(|s: &&str| !(*s).is_empty());

        let name = tokens.next();
        match name {
            Some(n) => {
                let mut params: Vec<String> = Vec::new();
                let mut hash: BTreeMap<String, Json> = BTreeMap::new();

                for t in tokens {
                    if t.contains_char('=') {
                        let kv = t.split('=').collect::<Vec<&str>>();
                        hash.insert(kv.get(0).unwrap().to_string(),
                                    kv.get(1).unwrap().parse::<Json>().unwrap());
                    } else {
                        params.push(t.to_string());
                    }
                }

                Ok(Helper{
                    name: n.to_string(),
                    params: params,
                    hash: hash,
                    template: Option::None,
                    inverse: Option::None,
                    block: block
                })
            },
            None =>
                Err(TemplateError)
        }
    }
}

impl Template {
    pub fn compile(source: String) -> Result<Template, TemplateError> {
        let mut helper_stack: RingBuf<Helper> = RingBuf::new();
        let mut template_stack: RingBuf<Template> = RingBuf::new();
        template_stack.push_front(Template{ elements: Vec::new() });

        let mut buffer: String = String::new();
        let mut state = ParserState::Text;

        let mut c:uint = 0;
        let source_len = source.chars().count();
        while c < source_len {
            let slice = source.slice_chars(c, min(c+3, source_len));
            state = match slice {
                "{{{" | "{{!" | "{{#" | "{{/" => {
                    c += 2;
                    if !buffer.is_empty() {
                        let mut t = template_stack.front_mut().unwrap();
                        t.elements.push(RawString(buffer.clone()));
                        buffer.clear();
                    }
                    match slice {
                        "{{{" => ParserState::HtmlExpression,
                        "{{!" => ParserState::Comment,
                        "{{#" => ParserState::HelperStart,
                        "{{/" => {
                            let t = template_stack.pop_front().unwrap();
                            let h = helper_stack.front_mut().unwrap();
                            if h.template.is_some() {
                                h.inverse = Some(t);
                            } else {
                                h.template = Some(t);
                            }
                            ParserState::HelperEnd
                        },
                        _ => ParserState::Invalid
                    }
                },
                "}}}" => {
                    c += 2;
                    let mut t = template_stack.front_mut().unwrap();
                    t.elements.push(HTMLExpression(
                        buffer.clone().trim_chars(' ').to_string()));
                    buffer.clear();
                    ParserState::Text
                },
                _ => {
                    match if slice.chars().count() > 2 { slice.slice_chars(0, 2) } else { slice } {
                        "{{" => {
                            c += 1;
                            if !buffer.is_empty() {
                                let mut t = template_stack.front_mut().unwrap();
                                t.elements.push(RawString(buffer.clone()));
                                buffer.clear();
                            }
                            ParserState::Expression
                        },
                        "}}" => {
                            c += 1;
                            match state {
                                ParserState::Expression => {
                                    if !buffer.is_empty() {
                                        // {{else}} within a helper block
                                        if buffer.trim_chars(' ') == "else" {
                                            buffer.clear(); // drop else
                                            let t = template_stack.pop_front().unwrap();
                                            let h = helper_stack.front_mut().unwrap();
                                            h.template = Some(t);
                                            template_stack.push_front(Template{ elements: Vec::new() });
                                            ParserState::Text
                                        } else if buffer.contains_char(' ') {
                                            //inline helper
                                            match Helper::parse(buffer.clone(), false){
                                                Ok(helper) => {
                                                    let mut t = template_stack.front_mut().unwrap();
                                                    t.elements.push(HelperExpression(helper));
                                                    buffer.clear();
                                                    ParserState::Text
                                                },
                                                Err(_) => ParserState::Invalid
                                            }
                                        } else {
                                            let mut t = template_stack.front_mut().unwrap();
                                            t.elements.push(Expression(
                                                buffer.clone().trim_chars(' ').to_string()));
                                            buffer.clear();
                                            ParserState::Text
                                        }

                                    } else {
                                        ParserState::Invalid
                                    }
                                },
                                ParserState::Comment => {
                                    let mut t = template_stack.front_mut().unwrap();
                                    t.elements.push(Comment(buffer.clone()));
                                    buffer.clear();
                                    ParserState::Text
                                },
                                ParserState::HelperStart => {
                                    match Helper::parse(buffer.clone(), true) {
                                        Ok(helper) => {
                                            helper_stack.push_front(helper);
                                            template_stack.push_front(Template{ elements: Vec::new() });

                                            buffer.clear();
                                            ParserState::Text
                                        },
                                        Err(_) => ParserState::Invalid
                                    }
                                },
                                ParserState::HelperEnd => {
                                    let name = buffer.trim_chars(' ').to_string();
                                    if name == helper_stack.front().unwrap().name {
                                        let h = helper_stack.pop_front().unwrap();
                                        let mut t = template_stack.front_mut().unwrap();
                                        t.elements.push(HelperBlock(h));
                                        buffer.clear();
                                        ParserState::Text
                                    } else {
                                        ParserState::Invalid
                                    }
                                },
                                _ => ParserState::Invalid
                            }
                        },
                        _ => {
                            buffer.push(slice.char_at(0));
                            state
                        }
                    }
                }
            };
            if state == ParserState::Invalid {
                return Err(TemplateError);
            }
            c += 1;
        }

        if !buffer.is_empty() {
            let mut t = template_stack.front_mut().unwrap();
            t.elements.push(TemplateElement::RawString(buffer.clone()));
        }

        if !helper_stack.is_empty() {
            return Err(TemplateError);
        }

        return Ok(template_stack.pop_front().unwrap());
    }
}

impl ToString for Template {
    fn to_string(&self) -> String {
        let mut buf = String::new();
        for v in self.elements.iter() {
            buf.push_str(v.to_string().as_slice());
        }
        buf
    }
}

#[deriving(PartialEq)]
pub enum TemplateElement {
    RawString(String),
    Expression(String),
    HTMLExpression(String),
    HelperExpression(Helper),
    HelperBlock(Helper),
    Comment(String),
}

impl ToString for TemplateElement {
    fn to_string(&self) -> String {
        match *self {
            RawString(ref v) => {
                v.clone()
            },
            Expression(ref v) => {
                // {{ escape to {
                format!("{{{{{}}}}}", v)
            },
            HTMLExpression(ref v) => {
                format!("{{{{{{{}}}}}}}", v)
            },
            HelperExpression(ref helper) => {
                helper.to_string()
            }
            HelperBlock(ref helper) => {
                helper.to_string()
            }
            Comment(ref v) => {
                format!("{{!{}}}", v)
            }
        }
    }
}

#[test]
fn test_parse_helper_start_tag() {
    let source = "if not name compare=1".to_string();
    let h = Helper::parse(source, true).ok().unwrap();

    assert_eq!(h.name, "if".to_string());
    assert_eq!(h.params, vec::<String>!["not".to_string(), "name".to_string()]);

    let key = "compare".to_string();
    let value = h.hash.get(&key).unwrap();
    assert_eq!(value.as_u64().unwrap(), 1u64);
}

#[test]
fn test_parse_template() {
    let source = "
{{title}} 你好
 {{{content}}}
{{#if date}}
good
{{else}}
bad
{{/if}}{{foo bar}}";
    let t = Template::compile(source.to_string()).ok().unwrap();

    assert_eq!(t.elements.len(), 8);
    assert_eq!((*t.elements.get(0).unwrap()).to_string(), "
".to_string());
    assert_eq!((*t.elements.get(1).unwrap()).to_string(),
               Expression("title".to_string()).to_string());

    assert_eq!((*t.elements.get(3).unwrap()).to_string(), "{{{content}}}".to_string());

    match *t.elements.get(5).unwrap() {
        HelperBlock(ref h) => {
            assert_eq!(h.name, "if".to_string());
            assert_eq!(h.params.len(), 1);
            assert_eq!(h.template.as_ref().unwrap().elements.len(), 1);
        },
        _ => {
            panic!("Helper expected here.");
        }
    };

    match *t.elements.get(7).unwrap() {
        HelperExpression(ref h) => {
            assert_eq!(h.name, "foo".to_string());
            assert_eq!(h.params.len(), 1);
            assert_eq!(*(h.params.get(0).unwrap()), "bar".to_string());
        },
        _ => {
            panic!("Helper expression here");
        }
    };
}

#[test]
fn test_helper_to_string() {
    let source = "{{#ifequals name compare=\"hello\"}}hello{{else}}good{{/ifequals}}".to_string();

    let t = Template::compile(source.to_string()).ok().unwrap();

    assert_eq!(t.elements.len(), 1);
    assert_eq!(t.elements.get(0).unwrap().to_string(), source);
}

#[test]
fn test_parse_error() {
    let source = "{{#ifequals name compare=\"hello\"}}hello{{else}}good".to_string();

    let t = Template::compile(source.to_string());

    assert!(t.is_err());
}
<commit_msg>rustup<commit_after>use std::cmp::min;
use std::iter::IteratorExt;
use std::collections::{BTreeMap, RingBuf};
use std::string::ToString;
use serialize::json::Json;

use self::TemplateElement::{RawString, Expression, HelperExpression,
                            HTMLExpression, HelperBlock, Comment};

#[deriving(PartialEq)]
pub struct Template {
    pub elements: Vec<TemplateElement>
}

#[deriving(PartialEq, Show)]
enum ParserState {
    Text,
    HtmlExpression,
    Comment,
    HelperStart,
    HelperEnd,
    Expression,
    Invalid
}

#[deriving(PartialEq)]
pub struct Helper {
    name: String,
    params: Vec<String>,
    hash: BTreeMap<String, Json>,
    template: Option<Template>,
    inverse: Option<Template>,
    block: bool
}

impl Helper {
    pub fn name(&self) -> &String {
        &self.name
    }

    pub fn params(&self) -> &Vec<String> {
        &self.params
    }

    pub fn hash(&self) -> &BTreeMap<String, Json> {
        &self.hash
    }

    pub fn template(&self) -> Option<&Template> {
        match self.template {
            Some(ref t) => {
                Some(t)
            },
            None => None
        }
    }

    pub fn inverse(&self) -> Option<&Template> {
        match self.inverse {
            Some(ref t) => {
                Some(t)
            },
            None => None
        }
    }

    pub fn is_block(&self) -> bool {
        self.block
    }
}

impl ToString for Helper {
    fn to_string(&self) -> String {

        let mut buf = String::new();

        if self.block {
            buf.push_str(format!("{{{{#{}", self.name).as_slice());
        } else {
            buf.push_str(format!("{{{{{}", self.name).as_slice());
        }

        for p in self.params.iter() {
            buf.push_str(format!(" {}", p).as_slice());
        }

        for k in self.hash.keys() {
            buf.push_str(format!(" {}={}", k, self.hash.get(k).unwrap()).as_slice());
        }

        buf.push_str("}}");

        if self.block {
            let tpl = self.template();
            if tpl.is_some() {
                buf.push_str(tpl.unwrap().to_string().as_slice());
            }
            let ivs = self.inverse();
            if ivs.is_some() {
                buf.push_str("{{else}}");
                buf.push_str(ivs.unwrap().to_string().as_slice());
            }
            buf.push_str(format!("{{{{/{}}}}}", self.name).as_slice());
        }
        buf
    }
}

#[deriving(Show, Copy)]
pub struct TemplateError;

impl Helper {
    fn parse(source: String, block: bool) -> Result<Helper, TemplateError> {
        let mut tokens = source.split(|&: c: char| c.is_whitespace())
            .filter(|s: &&str| !(*s).is_empty());

        let name = tokens.next();
        match name {
            Some(n) => {
                let mut params: Vec<String> = Vec::new();
                let mut hash: BTreeMap<String, Json> = BTreeMap::new();

                for t in tokens {
                    if t.contains_char('=') {
                        let kv = t.split('=').collect::<Vec<&str>>();
                        hash.insert(kv.get(0).unwrap().to_string(),
                                    kv.get(1).unwrap().parse::<Json>().unwrap());
                    } else {
                        params.push(t.to_string());
                    }
                }

                Ok(Helper{
                    name: n.to_string(),
                    params: params,
                    hash: hash,
                    template: Option::None,
                    inverse: Option::None,
                    block: block
                })
            },
            None =>
                Err(TemplateError)
        }
    }
}

impl Template {
    pub fn compile(source: String) -> Result<Template, TemplateError> {
        let mut helper_stack: RingBuf<Helper> = RingBuf::new();
        let mut template_stack: RingBuf<Template> = RingBuf::new();
        template_stack.push_front(Template{ elements: Vec::new() });

        let mut buffer: String = String::new();
        let mut state = ParserState::Text;

        let mut c:uint = 0;
        let source_len = source.chars().count();
        while c < source_len {
            let slice = source.slice_chars(c, min(c+3, source_len));
            state = match slice {
                "{{{" | "{{!" | "{{#" | "{{/" => {
                    c += 2;
                    if !buffer.is_empty() {
                        let mut t = template_stack.front_mut().unwrap();
                        t.elements.push(RawString(buffer.clone()));
                        buffer.clear();
                    }
                    match slice {
                        "{{{" => ParserState::HtmlExpression,
                        "{{!" => ParserState::Comment,
                        "{{#" => ParserState::HelperStart,
                        "{{/" => {
                            let t = template_stack.pop_front().unwrap();
                            let h = helper_stack.front_mut().unwrap();
                            if h.template.is_some() {
                                h.inverse = Some(t);
                            } else {
                                h.template = Some(t);
                            }
                            ParserState::HelperEnd
                        },
                        _ => ParserState::Invalid
                    }
                },
                "}}}" => {
                    c += 2;
                    let mut t = template_stack.front_mut().unwrap();
                    t.elements.push(HTMLExpression(
                        buffer.clone().trim_matches(' ').to_string()));
                    buffer.clear();
                    ParserState::Text
                },
                _ => {
                    match if slice.chars().count() > 2 { slice.slice_chars(0, 2) } else { slice } {
                        "{{" => {
                            c += 1;
                            if !buffer.is_empty() {
                                let mut t = template_stack.front_mut().unwrap();
                                t.elements.push(RawString(buffer.clone()));
                                buffer.clear();
                            }
                            ParserState::Expression
                        },
                        "}}" => {
                            c += 1;
                            match state {
                                ParserState::Expression => {
                                    if !buffer.is_empty() {
                                        // {{else}} within a helper block
                                        if buffer.trim_matches(' ') == "else" {
                                            buffer.clear(); // drop else
                                            let t = template_stack.pop_front().unwrap();
                                            let h = helper_stack.front_mut().unwrap();
                                            h.template = Some(t);
                                            template_stack.push_front(Template{ elements: Vec::new() });
                                            ParserState::Text
                                        } else if buffer.contains_char(' ') {
                                            //inline helper
                                            match Helper::parse(buffer.clone(), false){
                                                Ok(helper) => {
                                                    let mut t = template_stack.front_mut().unwrap();
                                                    t.elements.push(HelperExpression(helper));
                                                    buffer.clear();
                                                    ParserState::Text
                                                },
                                                Err(_) => ParserState::Invalid
                                            }
                                        } else {
                                            let mut t = template_stack.front_mut().unwrap();
                                            t.elements.push(Expression(
                                                buffer.clone().trim_matches(' ').to_string()));
                                            buffer.clear();
                                            ParserState::Text
                                        }

                                    } else {
                                        ParserState::Invalid
                                    }
                                },
                                ParserState::Comment => {
                                    let mut t = template_stack.front_mut().unwrap();
                                    t.elements.push(Comment(buffer.clone()));
                                    buffer.clear();
                                    ParserState::Text
                                },
                                ParserState::HelperStart => {
                                    match Helper::parse(buffer.clone(), true) {
                                        Ok(helper) => {
                                            helper_stack.push_front(helper);
                                            template_stack.push_front(Template{ elements: Vec::new() });

                                            buffer.clear();
                                            ParserState::Text
                                        },
                                        Err(_) => ParserState::Invalid
                                    }
                                },
                                ParserState::HelperEnd => {
                                    let name = buffer.trim_matches(' ').to_string();
                                    if name == helper_stack.front().unwrap().name {
                                        let h = helper_stack.pop_front().unwrap();
                                        let mut t = template_stack.front_mut().unwrap();
                                        t.elements.push(HelperBlock(h));
                                        buffer.clear();
                                        ParserState::Text
                                    } else {
                                        ParserState::Invalid
                                    }
                                },
                                _ => ParserState::Invalid
                            }
                        },
                        _ => {
                            buffer.push(slice.char_at(0));
                            state
                        }
                    }
                }
            };
            if state == ParserState::Invalid {
                return Err(TemplateError);
            }
            c += 1;
        }

        if !buffer.is_empty() {
            let mut t = template_stack.front_mut().unwrap();
            t.elements.push(TemplateElement::RawString(buffer.clone()));
        }

        if !helper_stack.is_empty() {
            return Err(TemplateError);
        }

        return Ok(template_stack.pop_front().unwrap());
    }
}

impl ToString for Template {
    fn to_string(&self) -> String {
        let mut buf = String::new();
        for v in self.elements.iter() {
            buf.push_str(v.to_string().as_slice());
        }
        buf
    }
}

#[deriving(PartialEq)]
pub enum TemplateElement {
    RawString(String),
    Expression(String),
    HTMLExpression(String),
    HelperExpression(Helper),
    HelperBlock(Helper),
    Comment(String),
}

impl ToString for TemplateElement {
    fn to_string(&self) -> String {
        match *self {
            RawString(ref v) => {
                v.clone()
            },
            Expression(ref v) => {
                // {{ escape to {
                format!("{{{{{}}}}}", v)
            },
            HTMLExpression(ref v) => {
                format!("{{{{{{{}}}}}}}", v)
            },
            HelperExpression(ref helper) => {
                helper.to_string()
            }
            HelperBlock(ref helper) => {
                helper.to_string()
            }
            Comment(ref v) => {
                format!("{{!{}}}", v)
            }
        }
    }
}

#[test]
fn test_parse_helper_start_tag() {
    let source = "if not name compare=1".to_string();
    let h = Helper::parse(source, true).ok().unwrap();

    assert_eq!(h.name, "if".to_string());
    assert_eq!(h.params, vec::<String>!["not".to_string(), "name".to_string()]);

    let key = "compare".to_string();
    let value = h.hash.get(&key).unwrap();
    assert_eq!(value.as_u64().unwrap(), 1u64);
}

#[test]
fn test_parse_template() {
    let source = "
{{title}} 你好
 {{{content}}}
{{#if date}}
good
{{else}}
bad
{{/if}}{{foo bar}}";
    let t = Template::compile(source.to_string()).ok().unwrap();

    assert_eq!(t.elements.len(), 8);
    assert_eq!((*t.elements.get(0).unwrap()).to_string(), "
".to_string());
    assert_eq!((*t.elements.get(1).unwrap()).to_string(),
               Expression("title".to_string()).to_string());

    assert_eq!((*t.elements.get(3).unwrap()).to_string(), "{{{content}}}".to_string());

    match *t.elements.get(5).unwrap() {
        HelperBlock(ref h) => {
            assert_eq!(h.name, "if".to_string());
            assert_eq!(h.params.len(), 1);
            assert_eq!(h.template.as_ref().unwrap().elements.len(), 1);
        },
        _ => {
            panic!("Helper expected here.");
        }
    };

    match *t.elements.get(7).unwrap() {
        HelperExpression(ref h) => {
            assert_eq!(h.name, "foo".to_string());
            assert_eq!(h.params.len(), 1);
            assert_eq!(*(h.params.get(0).unwrap()), "bar".to_string());
        },
        _ => {
            panic!("Helper expression here");
        }
    };
}

#[test]
fn test_helper_to_string() {
    let source = "{{#ifequals name compare=\"hello\"}}hello{{else}}good{{/ifequals}}".to_string();

    let t = Template::compile(source.to_string()).ok().unwrap();

    assert_eq!(t.elements.len(), 1);
    assert_eq!(t.elements.get(0).unwrap().to_string(), source);
}

#[test]
fn test_parse_error() {
    let source = "{{#ifequals name compare=\"hello\"}}hello{{else}}good".to_string();

    let t = Template::compile(source.to_string());

    assert!(t.is_err());
}
<|endoftext|>