<commit_before>use std::collections::HashMap;
use std::path::*;
use std::io::{Result, Write};
use case::CaseExt;

use env::Env;
use file_saver::*;
use library;
use nameutil::*;
use super::ffi_type::ffi_type;
use super::functions;
use super::statics;
use super::super::general::{self, tabs};
use traits::*;

pub fn generate(env: &Env) {
    println!("generating sys for {}", env.config.library_name);

    let path =  PathBuf::from(&env.config.target_path)
        .join(file_name_sys("lib"));
    println!("Generating file {:?}", path);

    save_to_file(path, &mut |w| generate_lib(w, env));
}

fn generate_lib<W: Write>(w: &mut W, env: &Env) -> Result<()>{
    try!(general::start_comments(w, &env.config));
    try!(statics::begin(w));

    try!(generate_extern_crates(w, env));
    try!(statics::after_extern_crates(w));

    if env.config.library_name != "GLib" {
        try!(statics::use_glib_ffi(w));
    } else {
        try!(statics::only_for_glib(w));
    }
    if env.config.library_name == "Gtk" {
        try!(statics::only_for_gtk(w));
    }

    let ns = env.library.namespace(library::MAIN_NAMESPACE);
    let classes = prepare(ns);
    let interfaces = prepare(ns);

    try!(generate_enums(w, &ns.name, &prepare(ns)));
    try!(generate_bitfields(w, &ns.name, &prepare(ns)));
    try!(functions::generate_callbacks(w, env, &prepare(ns)));
    try!(generate_records(w, env, &prepare_records(ns)));
    try!(generate_classes_structs(w, &classes));
    try!(generate_interfaces_structs(w, &interfaces));

    try!(writeln!(w, ""));
    try!(writeln!(w, "extern \"C\" {{"));
    try!(functions::generate_classes_funcs(w, env, &classes));
    try!(functions::generate_interfaces_funcs(w, env, &interfaces));

    //TODO: other functions
    try!(writeln!(w, "\n}}"));

    Ok(())
}

fn generate_extern_crates<W: Write>(w: &mut W, env: &Env) -> Result<()>{
    for library_name in &env.config.external_libraries {
        try!(writeln!(w, "extern crate {0}_sys as {0}_ffi;", crate_name(library_name)));
    }

    Ok(())
}

fn prepare<T: Ord>(ns: &library::Namespace) -> Vec<&T>
where library::Type: MaybeRef<T> {
    let mut vec: Vec<&T> = Vec::with_capacity(ns.types.len());
    for typ in ns.types.iter().filter_map(|t| t.as_ref()) {
        if let Some(ref x) = typ.maybe_ref() {
            vec.push(x);
        }
    }
    vec.sort();
    vec
}

fn prepare_records(ns: &library::Namespace) -> Vec<&library::Record> {
    let mut vec = Vec::with_capacity(ns.types.len());
    for typ in ns.types.iter().filter_map(|t| t.as_ref()) {
        if let Some(rec) = typ.maybe_ref_as::<library::Record>() {
            // We don't want the FooBarPrivate and similar records where FooBar is a type
            if ["Private", "Class", "Iface", "Interface"].iter()
                    .filter_map(|s| strip_suffix(&rec.name, s))
                    .any(|s| ns.index.get(s).is_some()) {
                continue;
            }
            vec.push(rec);
        }
    }
    vec.sort();
    vec
}

fn generate_bitfields<W: Write>(w: &mut W, ns_name: &str, items: &[&library::Bitfield])
        -> Result<()> {
    try!(writeln!(w, ""));
    for item in items {
        try!(writeln!(w, "bitflags! {{\n{}#[repr(C)]\n{0}flags {}: i32 {{", tabs(1), item.name));
        for member in &item.members {
            try!(writeln!(w, "{}const {} = {},",
                          tabs(2), strip_prefix(ns_name, &member.c_identifier), member.value));
        }
        try!(writeln!(w, "{}}}\n}}", tabs(1)));
        try!(writeln!(w, "pub type {} = {};", item.c_type, item.name));
        try!(writeln!(w, ""));
    }

    Ok(())
}

fn generate_enums<W: Write>(w: &mut W, ns_name: &str, items: &[&library::Enumeration])
        -> Result<()> {
    try!(writeln!(w, ""));
    for item in items {
        if item.members.len() == 1 {
            try!(writeln!(w, "pub type {} = i32;", item.name));
            try!(writeln!(w, "pub const {}: {} = {};",
                          item.members[0].c_identifier, item.name, item.members[0].value));
            try!(writeln!(w, "pub type {} = {};", item.c_type, item.name));
            try!(writeln!(w, ""));
            continue;
        }

        let mut vals = HashMap::new();
        try!(writeln!(w, "#[derive(Clone, Copy, Debug, Eq, PartialEq)]\n#[repr(C)]"));
        try!(writeln!(w, "pub enum {} {{", item.name));
        for member in &item.members {
            if vals.get(&member.value).is_some() {
                continue;
            }
            try!(writeln!(w, "{}{} = {},",
                          tabs(1), &prepare_enum_member_name(&member.name), member.value));
            vals.insert(member.value.clone(), member.name.clone());
        }
        try!(writeln!(w, "}}"));
        for member in &item.members {
            try!(writeln!(w, "pub const {}: {} = {1}::{};",
                          strip_prefix(ns_name, &member.c_identifier),
                          item.name, &prepare_enum_member_name(vals.get(&member.value).unwrap())));
        }
        try!(writeln!(w, "pub type {} = {};", item.c_type, item.name));
        try!(writeln!(w, ""));
    }

    Ok(())
}

fn prepare_enum_member_name(name: &str) -> String {
    let cameled = name.to_camel();
    if name.chars().next().unwrap().is_digit(10) {
        format!("_{}", cameled)
    } else {
        cameled
    }
}

fn generate_classes_structs<W: Write>(w: &mut W, classes: &[&library::Class]) -> Result<()> {
    try!(writeln!(w, ""));
    for klass in classes {
        try!(writeln!(w, "#[repr(C)]\npub struct {};", klass.c_type));
    }

    Ok(())
}

fn generate_interfaces_structs<W: Write>(w: &mut W, interfaces: &[&library::Interface]) -> Result<()> {
    try!(writeln!(w, ""));
    for interface in interfaces {
        try!(writeln!(w, "#[repr(C)]\npub struct {};", interface.c_type));
    }

    Ok(())
}

fn generate_records<W: Write>(w: &mut W, env: &Env, records: &[&library::Record]) -> Result<()> {
    try!(writeln!(w, ""));
    for record in records {
        let mut lines = Vec::new();
        let mut commented = false;
        for field in &record.fields {
            if env.library.type_(field.typ).maybe_ref_as::<library::Union>().is_some()
                    || field.bits.is_some() {
                warn!("Record `{}` field `{}` not expressible in Rust, truncated",
                      record.name, field.name);
                lines.push(format!("{}_truncated_record_marker: (),", tabs(1)));
                break;
            }
            let vis = if field.private { "" } else { "pub " };

            if let Some(ref c_type) = field.c_type {
                let name = mangle_keywords(&*field.name);
                let c_type = ffi_type(env, field.typ, c_type);
                lines.push(format!("{}{}{}: {},", tabs(1), vis, name, c_type.as_str()));
                if c_type.is_err() {
                    commented = true;
                }
            }
            else {
                let name = mangle_keywords(&*field.name);
                if let Some(ref func) =
                        env.library.type_(field.typ).maybe_ref_as::<library::Function>() {
                    let (com, sig) = functions::function_signature(env, func, true);
                    lines.push(format!("{}{}{}: fn{},", tabs(1), vis, name, sig));
                    commented |= com;
                }
                else {
                    lines.push(format!("{}{}{}: [{:?}],", tabs(1), vis, name, field.typ));
                    commented = true;
                }
            }
        }
        let comment = if commented { "//" } else { "" };
        if lines.is_empty() {
            try!(writeln!(w, "{}#[repr(C)]\n{0}pub struct {};\n", comment, record.c_type));
        }
        else {
            try!(writeln!(w, "{}#[repr(C)]\n{0}pub struct {} {{", comment, record.c_type));
            for line in lines {
                try!(writeln!(w, "{}{}", comment, line));
            }
            try!(writeln!(w, "{}}}\n", comment));
        }
    }
    Ok(())
}
<commit_msg>Do generate the FooClass, FooIface records<commit_after>use std::collections::HashMap;
use std::path::*;
use std::io::{Result, Write};
use case::CaseExt;

use env::Env;
use file_saver::*;
use library;
use nameutil::*;
use super::ffi_type::ffi_type;
use super::functions;
use super::statics;
use super::super::general::{self, tabs};
use traits::*;

pub fn generate(env: &Env) {
    println!("generating sys for {}", env.config.library_name);

    let path =  PathBuf::from(&env.config.target_path)
        .join(file_name_sys("lib"));
    println!("Generating file {:?}", path);

    save_to_file(path, &mut |w| generate_lib(w, env));
}

fn generate_lib<W: Write>(w: &mut W, env: &Env) -> Result<()>{
    try!(general::start_comments(w, &env.config));
    try!(statics::begin(w));

    try!(generate_extern_crates(w, env));
    try!(statics::after_extern_crates(w));

    if env.config.library_name != "GLib" {
        try!(statics::use_glib_ffi(w));
    } else {
        try!(statics::only_for_glib(w));
    }
    if env.config.library_name == "Gtk" {
        try!(statics::only_for_gtk(w));
    }

    let ns = env.library.namespace(library::MAIN_NAMESPACE);
    let classes = prepare(ns);
    let interfaces = prepare(ns);

    try!(generate_enums(w, &ns.name, &prepare(ns)));
    try!(generate_bitfields(w, &ns.name, &prepare(ns)));
    try!(functions::generate_callbacks(w, env, &prepare(ns)));
    try!(generate_records(w, env, &prepare_records(ns)));
    try!(generate_classes_structs(w, &classes));
    try!(generate_interfaces_structs(w, &interfaces));

    try!(writeln!(w, ""));
    try!(writeln!(w, "extern \"C\" {{"));
    try!(functions::generate_classes_funcs(w, env, &classes));
    try!(functions::generate_interfaces_funcs(w, env, &interfaces));

    //TODO: other functions
    try!(writeln!(w, "\n}}"));

    Ok(())
}

fn generate_extern_crates<W: Write>(w: &mut W, env: &Env) -> Result<()>{
    for library_name in &env.config.external_libraries {
        try!(writeln!(w, "extern crate {0}_sys as {0}_ffi;", crate_name(library_name)));
    }

    Ok(())
}

fn prepare<T: Ord>(ns: &library::Namespace) -> Vec<&T>
where library::Type: MaybeRef<T> {
    let mut vec: Vec<&T> = Vec::with_capacity(ns.types.len());
    for typ in ns.types.iter().filter_map(|t| t.as_ref()) {
        if let Some(ref x) = typ.maybe_ref() {
            vec.push(x);
        }
    }
    vec.sort();
    vec
}

fn prepare_records(ns: &library::Namespace) -> Vec<&library::Record> {
    let mut vec = Vec::with_capacity(ns.types.len());
    for typ in ns.types.iter().filter_map(|t| t.as_ref()) {
        if let Some(rec) = typ.maybe_ref_as::<library::Record>() {
            // We don't want the FooBarPrivate and similar records where FooBar is a type
            if ["Private"].iter()
                    .filter_map(|s| strip_suffix(&rec.name, s))
                    .any(|s| ns.index.get(s).is_some()) {
                continue;
            }
            vec.push(rec);
        }
    }
    vec.sort();
    vec
}

fn generate_bitfields<W: Write>(w: &mut W, ns_name: &str, items: &[&library::Bitfield])
        -> Result<()> {
    try!(writeln!(w, ""));
    for item in items {
        try!(writeln!(w, "bitflags! {{\n{}#[repr(C)]\n{0}flags {}: i32 {{", tabs(1), item.name));
        for member in &item.members {
            try!(writeln!(w, "{}const {} = {},",
                          tabs(2), strip_prefix(ns_name, &member.c_identifier), member.value));
        }
        try!(writeln!(w, "{}}}\n}}", tabs(1)));
        try!(writeln!(w, "pub type {} = {};", item.c_type, item.name));
        try!(writeln!(w, ""));
    }

    Ok(())
}

fn generate_enums<W: Write>(w: &mut W, ns_name: &str, items: &[&library::Enumeration])
        -> Result<()> {
    try!(writeln!(w, ""));
    for item in items {
        if item.members.len() == 1 {
            try!(writeln!(w, "pub type {} = i32;", item.name));
            try!(writeln!(w, "pub const {}: {} = {};",
                          item.members[0].c_identifier, item.name, item.members[0].value));
            try!(writeln!(w, "pub type {} = {};", item.c_type, item.name));
            try!(writeln!(w, ""));
            continue;
        }

        let mut vals = HashMap::new();
        try!(writeln!(w, "#[derive(Clone, Copy, Debug, Eq, PartialEq)]\n#[repr(C)]"));
        try!(writeln!(w, "pub enum {} {{", item.name));
        for member in &item.members {
            if vals.get(&member.value).is_some() {
                continue;
            }
            try!(writeln!(w, "{}{} = {},",
                          tabs(1), &prepare_enum_member_name(&member.name), member.value));
            vals.insert(member.value.clone(), member.name.clone());
        }
        try!(writeln!(w, "}}"));
        for member in &item.members {
            try!(writeln!(w, "pub const {}: {} = {1}::{};",
                          strip_prefix(ns_name, &member.c_identifier),
                          item.name, &prepare_enum_member_name(vals.get(&member.value).unwrap())));
        }
        try!(writeln!(w, "pub type {} = {};", item.c_type, item.name));
        try!(writeln!(w, ""));
    }

    Ok(())
}

fn prepare_enum_member_name(name: &str) -> String {
    let cameled = name.to_camel();
    if name.chars().next().unwrap().is_digit(10) {
        format!("_{}", cameled)
    } else {
        cameled
    }
}

fn generate_classes_structs<W: Write>(w: &mut W, classes: &[&library::Class]) -> Result<()> {
    try!(writeln!(w, ""));
    for klass in classes {
        try!(writeln!(w, "#[repr(C)]\npub struct {};", klass.c_type));
    }

    Ok(())
}

fn generate_interfaces_structs<W: Write>(w: &mut W, interfaces: &[&library::Interface]) -> Result<()> {
    try!(writeln!(w, ""));
    for interface in interfaces {
        try!(writeln!(w, "#[repr(C)]\npub struct {};", interface.c_type));
    }

    Ok(())
}

fn generate_records<W: Write>(w: &mut W, env: &Env, records: &[&library::Record]) -> Result<()> {
    try!(writeln!(w, ""));
    for record in records {
        let mut lines = Vec::new();
        let mut commented = false;
        for field in &record.fields {
            if env.library.type_(field.typ).maybe_ref_as::<library::Union>().is_some()
                    || field.bits.is_some() {
                warn!("Record `{}` field `{}` not expressible in Rust, truncated",
                      record.name, field.name);
                lines.push(format!("{}_truncated_record_marker: (),", tabs(1)));
                break;
            }
            let vis = if field.private { "" } else { "pub " };

            if let Some(ref c_type) = field.c_type {
                let name = mangle_keywords(&*field.name);
                let c_type = ffi_type(env, field.typ, c_type);
                lines.push(format!("{}{}{}: {},", tabs(1), vis, name, c_type.as_str()));
                if c_type.is_err() {
                    commented = true;
                }
            }
            else {
                let name = mangle_keywords(&*field.name);
                if let Some(ref func) =
                        env.library.type_(field.typ).maybe_ref_as::<library::Function>() {
                    let (com, sig) = functions::function_signature(env, func, true);
                    lines.push(format!("{}{}{}: fn{},", tabs(1), vis, name, sig));
                    commented |= com;
                }
                else {
                    lines.push(format!("{}{}{}: [{:?}],", tabs(1), vis, name, field.typ));
                    commented = true;
                }
            }
        }
        let comment = if commented { "//" } else { "" };
        if lines.is_empty() {
            try!(writeln!(w, "{}#[repr(C)]\n{0}pub struct {};\n", comment, record.c_type));
        }
        else {
            try!(writeln!(w, "{}#[repr(C)]\n{0}pub struct {} {{", comment, record.c_type));
            for line in lines {
                try!(writeln!(w, "{}{}", comment, line));
            }
            try!(writeln!(w, "{}}}\n", comment));
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>use std::sync::{Arc, RwLock};
use std::option::Option;
use std::clone::Clone;

use uuid::Uuid;

type OptInt32 = Option<i32>;
type OptInt64 = Option<i64>;
type OptString = Option<String>;

#[derive(Clone)]
pub enum CellType {
    Int32 (OptInt32),
    Int64 (OptInt64),
    Text  (OptString)
}

pub struct CellFormat {
    pub fieldname: String,
    pub fieldtype: CellType
}

pub struct RecFormat {
    pub cols: Vec<CellFormat>
}

impl RecFormat {
    pub fn new(fields: Vec<CellFormat>) -> Arc<RwLock<RecFormat>> {
        return Arc::new(RwLock::new(RecFormat { cols: fields } ));
    }
}

trait Formattable {
    pub fn own_format(&mut self) {
        self.format = RecFormat::new(self.format.cols);
    }
}

pub struct Rec {
    pub guid: Uuid,
    pub cells: Vec<CellType>,
    pub format: Arc<RwLock<RecFormat>>
}

impl Rec {
    pub fn new(format: Arc<RwLock<RecFormat>>) -> Rec {
        let cells = format
                        .read()
                        .unwrap()
                        .cols
                        .iter()
                        .map(|x| x.type_.clone()).collect::<Vec<CellType>>();
        return Rec { guid: Uuid::new_v4(),
                     cells: cells,
                     format: format
                   };
    }
}

impl Formattable for Rec {
}

pub struct Table {
    pub name: String,
    pub rows: Vec<Rec>,
    pub format: Arc<RwLock<RecFormat>>
}

impl Table {
    pub fn new(format: Arc<RwLock<RecFormat>>) -> Table {
        return Table {
            format: format
        };
    }
}
<commit_msg>Remove pub<commit_after>use std::sync::{Arc, RwLock};
use std::option::Option;
use std::clone::Clone;

use uuid::Uuid;

type OptInt32 = Option<i32>;
type OptInt64 = Option<i64>;
type OptString = Option<String>;

#[derive(Clone)]
pub enum CellType {
    Int32 (OptInt32),
    Int64 (OptInt64),
    Text  (OptString)
}

pub struct CellFormat {
    pub fieldname: String,
    pub fieldtype: CellType
}

pub struct RecFormat {
    pub cols: Vec<CellFormat>
}

impl RecFormat {
    pub fn new(fields: Vec<CellFormat>) -> Arc<RwLock<RecFormat>> {
        return Arc::new(RwLock::new(RecFormat { cols: fields } ));
    }
}

trait Formattable {
    fn own_format(&mut self) {
        self.format = RecFormat::new(self.format.cols);
    }
}

pub struct Rec {
    pub guid: Uuid,
    pub cells: Vec<CellType>,
    pub format: Arc<RwLock<RecFormat>>
}

impl Rec {
    pub fn new(format: Arc<RwLock<RecFormat>>) -> Rec {
        let cells = format
                        .read()
                        .unwrap()
                        .cols
                        .iter()
                        .map(|x| x.type_.clone()).collect::<Vec<CellType>>();
        return Rec { guid: Uuid::new_v4(),
                     cells: cells,
                     format: format
                   };
    }
}

impl Formattable for Rec {
}

pub struct Table {
    pub name: String,
    pub rows: Vec<Rec>,
    pub format: Arc<RwLock<RecFormat>>
}

impl Table {
    pub fn new(format: Arc<RwLock<RecFormat>>) -> Table {
        return Table {
            format: format
        };
    }
}
<|endoftext|>


<commit_before>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

pub extern mod core_foundation;
pub extern mod io_surface;
pub extern mod opengles;

use mod platform::opengles::gl2;

use core::cast::transmute;
use core::ptr::{null, to_unsafe_ptr};
use geom::size::Size2D;
use platform::io_surface::{IOSurface, kIOSurfaceBytesPerElement, kIOSurfaceBytesPerRow};
use platform::io_surface::{kIOSurfaceHeight, kIOSurfaceIsGlobal, kIOSurfaceWidth};
use platform::io_surface::IOSurfaceMethods;
use platform::opengles::cgl::{CGLChoosePixelFormat, CGLContextObj, CGLCreateContext};
use platform::opengles::cgl::{CGLLockContext, CGLSetCurrentContext, CGLTexImageIOSurface2D};
use platform::opengles::cgl::{kCGLNoError, kCGLPFACompliant, kCGLPFADoubleBuffer};
use platform::opengles::gl2::{BGRA, CLAMP_TO_EDGE, COLOR_ATTACHMENT0, FRAMEBUFFER};
use platform::opengles::gl2::{FRAMEBUFFER_COMPLETE, GLenum, GLint, GLsizei, GLuint, LINEAR};
use platform::opengles::gl2::{NEAREST, RGBA, TEXTURE_MAG_FILTER, TEXTURE_MIN_FILTER};
use platform::opengles::gl2::{TEXTURE_RECTANGLE_ARB, TEXTURE_WRAP_S, TEXTURE_WRAP_T};
use platform::opengles::gl2::{UNSIGNED_INT_8_8_8_8_REV};

use base::ShareContext;

// FIXME: This is not good.
#[link_args="-framework IOSurface -framework CoreFoundation"]
#[nolink]
extern {}

pub type Context = MacContext;

pub struct MacContext {
    surface: IOSurface,
    framebuffer: GLuint,
    texture: GLuint

    // FIXME: Needs drop.
}

pub fn init_cgl() -> CGLContextObj {
    // FIXME: We should expose some higher-level, safe abstractions inside the CGL module.
    unsafe {
        // Choose a pixel format.
        let attributes = [ kCGLPFADoubleBuffer, kCGLPFACompliant, 0 ];
        let pixel_format = null();
        let pixel_format_count = 1;
        let gl_error = CGLChoosePixelFormat(transmute(&attributes[0]),
                                            to_unsafe_ptr(&pixel_format),
                                            to_unsafe_ptr(&pixel_format_count));
        assert!(gl_error == kCGLNoError);

        // Create the context.
        let cgl_context = null();
        let gl_error = CGLCreateContext(pixel_format, null(), to_unsafe_ptr(&cgl_context));
        assert!(gl_error == kCGLNoError);

        // Set the context.
        let gl_error = CGLSetCurrentContext(cgl_context);
        assert!(gl_error == kCGLNoError);

        // Lock the context.
        let gl_error = CGLLockContext(cgl_context);
        assert!(gl_error == kCGLNoError);

        return cgl_context;
    }
}

pub fn init_surface(size: Size2D<int>) -> IOSurface {
    use platform::core_foundation::boolean::CFBoolean;
    use number = platform::core_foundation::number::CFNumber::new;
    use string = platform::core_foundation::string::CFString::wrap_extern;

    // TODO: dictionary constructor should be less ridiculous.
    // Or, we could add bindings for mutable dictionaries.
    let k_width = string(kIOSurfaceWidth);
    let v_width = number(size.width as i32);

    let k_height = string(kIOSurfaceHeight);
    let v_height = number(size.height as i32);

    let k_bytes_per_row = string(kIOSurfaceBytesPerRow);
    let v_bytes_per_row = number(size.width as i32 * 4);

    let k_bytes_per_elem = string(kIOSurfaceBytesPerElement);
    let v_bytes_per_elem = number(4i32);

    let k_is_global = string(kIOSurfaceIsGlobal);
    let v_is_global = CFBoolean::true_value();

    io_surface::new(&core_foundation::dictionary::CFDictionary::new([
        (*k_width.borrow_ref(),          *v_width.borrow_type_ref()),
        (*k_height.borrow_ref(),         *v_height.borrow_type_ref()),
        (*k_bytes_per_row.borrow_ref(),  *v_bytes_per_row.borrow_type_ref()),
        (*k_bytes_per_elem.borrow_ref(), *v_bytes_per_elem.borrow_type_ref()),
        (*k_is_global.borrow_ref(),      *v_is_global.borrow_type_ref()),
    ]))
}

pub fn init_texture() -> GLuint {
    gl2::enable(TEXTURE_RECTANGLE_ARB);

    let texture = gl2::gen_textures(1)[0];
    gl2::bind_texture(TEXTURE_RECTANGLE_ARB, texture);
    gl2::tex_parameter_i(TEXTURE_RECTANGLE_ARB, TEXTURE_WRAP_S, CLAMP_TO_EDGE as GLint);
    gl2::tex_parameter_i(TEXTURE_RECTANGLE_ARB, TEXTURE_WRAP_T, CLAMP_TO_EDGE as GLint);
    gl2::tex_parameter_i(TEXTURE_RECTANGLE_ARB, TEXTURE_MAG_FILTER, LINEAR as GLint);
    gl2::tex_parameter_i(TEXTURE_RECTANGLE_ARB, TEXTURE_MIN_FILTER, NEAREST as GLint);
    return texture;
}

// Assumes the texture is already bound via gl2::bind_texture().
pub fn bind_surface_to_texture(context: &CGLContextObj, surface: &IOSurface, size: Size2D<int>) {
    // FIXME: There should be safe wrappers for this.
    unsafe {
        let gl_error = CGLTexImageIOSurface2D(*context,
                                              TEXTURE_RECTANGLE_ARB,
                                              RGBA as GLenum,
                                              size.width as GLsizei,
                                              size.height as GLsizei,
                                              BGRA as GLenum,
                                              UNSIGNED_INT_8_8_8_8_REV,
                                              transmute(copy surface.obj),
                                              0);
        assert!(gl_error == kCGLNoError);
    }
}

pub fn bind_texture_to_framebuffer(texture: GLuint) {
    gl2::bind_texture(TEXTURE_RECTANGLE_ARB, 0);
    gl2::framebuffer_texture_2d(FRAMEBUFFER, COLOR_ATTACHMENT0, TEXTURE_RECTANGLE_ARB, texture, 0);
    assert!(gl2::check_framebuffer_status(FRAMEBUFFER) == FRAMEBUFFER_COMPLETE);
}

impl ShareContext for MacContext {
    fn new(size: Size2D<int>) -> MacContext {
        // Initialize CGL.
        let context = init_cgl();

        // Create the surface.
        let surface = init_surface(copy size);

        // Create a framebuffer.
        let framebuffer = gl2::gen_framebuffers(1)[0];
        gl2::bind_framebuffer(FRAMEBUFFER, framebuffer);

        // Create and bind to the texture.
        let texture = init_texture();
        bind_surface_to_texture(&context, &surface, size);

        // Bind the texture to the framebuffer.
        bind_texture_to_framebuffer(texture);

        MacContext {
            surface: surface,
            framebuffer: framebuffer,
            texture: texture
        }
    }

    fn flush(&self) {
        gl2::finish();
    }

    fn id(&self) -> int {
        self.surface.get_id() as int
    }
}

<commit_msg>Fix build error.<commit_after>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

pub extern mod core_foundation;
pub extern mod io_surface;
pub extern mod opengles;

use platform::opengles::gl2;

use core::cast::transmute;
use core::ptr::{null, to_unsafe_ptr};
use geom::size::Size2D;
use platform::io_surface::{IOSurface, kIOSurfaceBytesPerElement, kIOSurfaceBytesPerRow};
use platform::io_surface::{kIOSurfaceHeight, kIOSurfaceIsGlobal, kIOSurfaceWidth};
use platform::io_surface::IOSurfaceMethods;
use platform::opengles::cgl::{CGLChoosePixelFormat, CGLContextObj, CGLCreateContext};
use platform::opengles::cgl::{CGLLockContext, CGLSetCurrentContext, CGLTexImageIOSurface2D};
use platform::opengles::cgl::{kCGLNoError, kCGLPFACompliant, kCGLPFADoubleBuffer};
use platform::opengles::gl2::{BGRA, CLAMP_TO_EDGE, COLOR_ATTACHMENT0, FRAMEBUFFER};
use platform::opengles::gl2::{FRAMEBUFFER_COMPLETE, GLenum, GLint, GLsizei, GLuint, LINEAR};
use platform::opengles::gl2::{NEAREST, RGBA, TEXTURE_MAG_FILTER, TEXTURE_MIN_FILTER};
use platform::opengles::gl2::{TEXTURE_RECTANGLE_ARB, TEXTURE_WRAP_S, TEXTURE_WRAP_T};
use platform::opengles::gl2::{UNSIGNED_INT_8_8_8_8_REV};

use base::ShareContext;

// FIXME: This is not good.
#[link_args="-framework IOSurface -framework CoreFoundation"]
#[nolink]
extern {}

pub type Context = MacContext;

pub struct MacContext {
    surface: IOSurface,
    framebuffer: GLuint,
    texture: GLuint

    // FIXME: Needs drop.
}

pub fn init_cgl() -> CGLContextObj {
    // FIXME: We should expose some higher-level, safe abstractions inside the CGL module.
    unsafe {
        // Choose a pixel format.
        let attributes = [ kCGLPFADoubleBuffer, kCGLPFACompliant, 0 ];
        let pixel_format = null();
        let pixel_format_count = 1;
        let gl_error = CGLChoosePixelFormat(transmute(&attributes[0]),
                                            to_unsafe_ptr(&pixel_format),
                                            to_unsafe_ptr(&pixel_format_count));
        assert!(gl_error == kCGLNoError);

        // Create the context.
        let cgl_context = null();
        let gl_error = CGLCreateContext(pixel_format, null(), to_unsafe_ptr(&cgl_context));
        assert!(gl_error == kCGLNoError);

        // Set the context.
        let gl_error = CGLSetCurrentContext(cgl_context);
        assert!(gl_error == kCGLNoError);

        // Lock the context.
        let gl_error = CGLLockContext(cgl_context);
        assert!(gl_error == kCGLNoError);

        return cgl_context;
    }
}

pub fn init_surface(size: Size2D<int>) -> IOSurface {
    use platform::core_foundation::boolean::CFBoolean;
    use number = platform::core_foundation::number::CFNumber::new;
    use string = platform::core_foundation::string::CFString::wrap_extern;

    // TODO: dictionary constructor should be less ridiculous.
    // Or, we could add bindings for mutable dictionaries.
    let k_width = string(kIOSurfaceWidth);
    let v_width = number(size.width as i32);

    let k_height = string(kIOSurfaceHeight);
    let v_height = number(size.height as i32);

    let k_bytes_per_row = string(kIOSurfaceBytesPerRow);
    let v_bytes_per_row = number(size.width as i32 * 4);

    let k_bytes_per_elem = string(kIOSurfaceBytesPerElement);
    let v_bytes_per_elem = number(4i32);

    let k_is_global = string(kIOSurfaceIsGlobal);
    let v_is_global = CFBoolean::true_value();

    io_surface::new(&core_foundation::dictionary::CFDictionary::new([
        (*k_width.borrow_ref(),          *v_width.borrow_type_ref()),
        (*k_height.borrow_ref(),         *v_height.borrow_type_ref()),
        (*k_bytes_per_row.borrow_ref(),  *v_bytes_per_row.borrow_type_ref()),
        (*k_bytes_per_elem.borrow_ref(), *v_bytes_per_elem.borrow_type_ref()),
        (*k_is_global.borrow_ref(),      *v_is_global.borrow_type_ref()),
    ]))
}

pub fn init_texture() -> GLuint {
    gl2::enable(TEXTURE_RECTANGLE_ARB);

    let texture = gl2::gen_textures(1)[0];
    gl2::bind_texture(TEXTURE_RECTANGLE_ARB, texture);
    gl2::tex_parameter_i(TEXTURE_RECTANGLE_ARB, TEXTURE_WRAP_S, CLAMP_TO_EDGE as GLint);
    gl2::tex_parameter_i(TEXTURE_RECTANGLE_ARB, TEXTURE_WRAP_T, CLAMP_TO_EDGE as GLint);
    gl2::tex_parameter_i(TEXTURE_RECTANGLE_ARB, TEXTURE_MAG_FILTER, LINEAR as GLint);
    gl2::tex_parameter_i(TEXTURE_RECTANGLE_ARB, TEXTURE_MIN_FILTER, NEAREST as GLint);
    return texture;
}

// Assumes the texture is already bound via gl2::bind_texture().
pub fn bind_surface_to_texture(context: &CGLContextObj, surface: &IOSurface, size: Size2D<int>) {
    // FIXME: There should be safe wrappers for this.
    unsafe {
        let gl_error = CGLTexImageIOSurface2D(*context,
                                              TEXTURE_RECTANGLE_ARB,
                                              RGBA as GLenum,
                                              size.width as GLsizei,
                                              size.height as GLsizei,
                                              BGRA as GLenum,
                                              UNSIGNED_INT_8_8_8_8_REV,
                                              transmute(copy surface.obj),
                                              0);
        assert!(gl_error == kCGLNoError);
    }
}

pub fn bind_texture_to_framebuffer(texture: GLuint) {
    gl2::bind_texture(TEXTURE_RECTANGLE_ARB, 0);
    gl2::framebuffer_texture_2d(FRAMEBUFFER, COLOR_ATTACHMENT0, TEXTURE_RECTANGLE_ARB, texture, 0);
    assert!(gl2::check_framebuffer_status(FRAMEBUFFER) == FRAMEBUFFER_COMPLETE);
}

impl ShareContext for MacContext {
    fn new(size: Size2D<int>) -> MacContext {
        // Initialize CGL.
        let context = init_cgl();

        // Create the surface.
        let surface = init_surface(copy size);

        // Create a framebuffer.
        let framebuffer = gl2::gen_framebuffers(1)[0];
        gl2::bind_framebuffer(FRAMEBUFFER, framebuffer);

        // Create and bind to the texture.
        let texture = init_texture();
        bind_surface_to_texture(&context, &surface, size);

        // Bind the texture to the framebuffer.
        bind_texture_to_framebuffer(texture);

        MacContext {
            surface: surface,
            framebuffer: framebuffer,
            texture: texture
        }
    }

    fn flush(&self) {
        gl2::finish();
    }

    fn id(&self) -> int {
        self.surface.get_id() as int
    }
}

<|endoftext|>


<commit_before>use rustc_data_structures::fx::FxHashMap;
use rustc_hir::def_id::DefId;
use rustc_infer::infer::InferCtxt;
use rustc_middle::ty::{self, TyCtxt, TypeFoldable};
use rustc_span::Span;
use rustc_trait_selection::opaque_types::InferCtxtExt;

use super::RegionInferenceContext;

impl<'tcx> RegionInferenceContext<'tcx> {
    /// Resolve any opaque types that were encountered while borrow checking
    /// this item. This is then used to get the type in the `type_of` query.
    ///
    /// For example consider `fn f<'a>(x: &'a i32) -> impl Sized + 'a { x }`.
    /// This is lowered to give HIR something like
    ///
    /// type f<'a>::_Return<'_a> = impl Sized + '_a;
    /// fn f<'a>(x: &'a i32) -> f<'static>::_Return<'a> { x }
    ///
    /// When checking the return type record the type from the return and the
    /// type used in the return value. In this case they might be `_Return<'1>`
    /// and `&'2 i32` respectively.
    ///
    /// Once we to this method, we have completed region inference and want to
    /// call `infer_opaque_definition_from_instantiation` to get the inferred
    /// type of `_Return<'_a>`. `infer_opaque_definition_from_instantiation`
    /// compares lifetimes directly, so we need to map the inference variables
    /// back to concrete lifetimes: `'static`, `ReEarlyBound` or `ReFree`.
    ///
    /// First we map all the lifetimes in the concrete type to an equal
    /// universal region that occurs in the concrete type's substs, in this case
    /// this would result in `&'1 i32`. We only consider regions in the substs
    /// in case there is an equal region that does not. For example, this should
    /// be allowed:
    /// `fn f<'a: 'b, 'b: 'a>(x: *mut &'b i32) -> impl Sized + 'a { x }`
    ///
    /// Then we map the regions in both the type and the subst to their
    /// `external_name` giving `concrete_type = &'a i32`,
    /// `substs = ['static, 'a]`. This will then allow
    /// `infer_opaque_definition_from_instantiation` to determine that
    /// `_Return<'_a> = &'_a i32`.
    ///
    /// There's a slight complication around closures. Given
    /// `fn f<'a: 'a>() { || {} }` the closure's type is something like
    /// `f::<'a>::{{closure}}`. The region parameter from f is essentially
    /// ignored by type checking so ends up being inferred to an empty region.
    /// Calling `universal_upper_bound` for such a region gives `fr_fn_body`,
    /// which has no `external_name` in which case we use `'empty` as the
    /// region to pass to `infer_opaque_definition_from_instantiation`.
    pub(in crate::borrow_check) fn infer_opaque_types(
        &self,
        infcx: &InferCtxt<'_, 'tcx>,
        opaque_ty_decls: FxHashMap<DefId, ty::ResolvedOpaqueTy<'tcx>>,
        span: Span,
    ) -> FxHashMap<DefId, ty::ResolvedOpaqueTy<'tcx>> {
        opaque_ty_decls
            .into_iter()
            .map(|(opaque_def_id, ty::ResolvedOpaqueTy { concrete_type, substs })| {
                debug!(
                    "infer_opaque_types(concrete_type = {:?}, substs = {:?})",
                    concrete_type, substs
                );

                let mut subst_regions = vec![self.universal_regions.fr_static];
                let universal_substs =
                    infcx.tcx.fold_regions(substs, &mut false, |region, _| match *region {
                        ty::ReVar(vid) => {
                            subst_regions.push(vid);
                            self.definitions[vid].external_name.unwrap_or_else(|| {
                                infcx.tcx.sess.delay_span_bug(
                                    span,
                                    "opaque type with non-universal region substs",
                                );
                                infcx.tcx.lifetimes.re_static
                            })
                        }
                        // We don't fold regions in the predicates of opaque
                        // types to `ReVar`s. This means that in a case like
                        //
                        // fn f<'a: 'a>() -> impl Iterator<Item = impl Sized>
                        //
                        // The inner opaque type has `'static` in its substs.
                        ty::ReStatic => region,
                        _ => {
                            infcx.tcx.sess.delay_span_bug(
                                span,
                                &format!("unexpected concrete region in borrowck: {:?}", region),
                            );
                            region
                        }
                    });

                subst_regions.sort();
                subst_regions.dedup();

                let universal_concrete_type =
                    infcx.tcx.fold_regions(concrete_type, &mut false, |region, _| match *region {
                        ty::ReVar(vid) => subst_regions
                            .iter()
                            .find(|ur_vid| self.eval_equal(vid, **ur_vid))
                            .and_then(|ur_vid| self.definitions[*ur_vid].external_name)
                            .unwrap_or(infcx.tcx.lifetimes.re_root_empty),
                        ty::ReLateBound(..) => region,
                        _ => {
                            infcx.tcx.sess.delay_span_bug(
                                span,
                                &format!("unexpected concrete region in borrowck: {:?}", region),
                            );
                            region
                        }
                    });

                debug!(
                    "infer_opaque_types(universal_concrete_type = {:?}, universal_substs = {:?})",
                    universal_concrete_type, universal_substs
                );

                let remapped_type = infcx.infer_opaque_definition_from_instantiation(
                    opaque_def_id,
                    universal_substs,
                    universal_concrete_type,
                    span,
                );
                (
                    opaque_def_id,
                    ty::ResolvedOpaqueTy { concrete_type: remapped_type, substs: universal_substs },
                )
            })
            .collect()
    }

    /// Map the regions in the type to named regions. This is similar to what
    /// `infer_opaque_types` does, but can infer any universal region, not only
    /// ones from the substs for the opaque type. It also doesn't double check
    /// that the regions produced are in fact equal to the named region they are
    /// replaced with. This is fine because this function is only to improve the
    /// region names in error messages.
    pub(in crate::borrow_check) fn name_regions<T>(&self, tcx: TyCtxt<'tcx>, ty: T) -> T
    where
        T: TypeFoldable<'tcx>,
    {
        tcx.fold_regions(ty, &mut false, |region, _| match *region {
            ty::ReVar(vid) => {
                // Find something that we can name
                let upper_bound = self.approx_universal_upper_bound(vid);
                self.definitions[upper_bound].external_name.unwrap_or(region)
            }
            _ => region,
        })
    }
}
<commit_msg>Use tracing instrumentation for better bug diagnosing<commit_after>use rustc_data_structures::fx::FxHashMap;
use rustc_hir::def_id::DefId;
use rustc_infer::infer::InferCtxt;
use rustc_middle::ty::{self, TyCtxt, TypeFoldable};
use rustc_span::Span;
use rustc_trait_selection::opaque_types::InferCtxtExt;

use super::RegionInferenceContext;

impl<'tcx> RegionInferenceContext<'tcx> {
    /// Resolve any opaque types that were encountered while borrow checking
    /// this item. This is then used to get the type in the `type_of` query.
    ///
    /// For example consider `fn f<'a>(x: &'a i32) -> impl Sized + 'a { x }`.
    /// This is lowered to give HIR something like
    ///
    /// type f<'a>::_Return<'_a> = impl Sized + '_a;
    /// fn f<'a>(x: &'a i32) -> f<'static>::_Return<'a> { x }
    ///
    /// When checking the return type record the type from the return and the
    /// type used in the return value. In this case they might be `_Return<'1>`
    /// and `&'2 i32` respectively.
    ///
    /// Once we to this method, we have completed region inference and want to
    /// call `infer_opaque_definition_from_instantiation` to get the inferred
    /// type of `_Return<'_a>`. `infer_opaque_definition_from_instantiation`
    /// compares lifetimes directly, so we need to map the inference variables
    /// back to concrete lifetimes: `'static`, `ReEarlyBound` or `ReFree`.
    ///
    /// First we map all the lifetimes in the concrete type to an equal
    /// universal region that occurs in the concrete type's substs, in this case
    /// this would result in `&'1 i32`. We only consider regions in the substs
    /// in case there is an equal region that does not. For example, this should
    /// be allowed:
    /// `fn f<'a: 'b, 'b: 'a>(x: *mut &'b i32) -> impl Sized + 'a { x }`
    ///
    /// Then we map the regions in both the type and the subst to their
    /// `external_name` giving `concrete_type = &'a i32`,
    /// `substs = ['static, 'a]`. This will then allow
    /// `infer_opaque_definition_from_instantiation` to determine that
    /// `_Return<'_a> = &'_a i32`.
    ///
    /// There's a slight complication around closures. Given
    /// `fn f<'a: 'a>() { || {} }` the closure's type is something like
    /// `f::<'a>::{{closure}}`. The region parameter from f is essentially
    /// ignored by type checking so ends up being inferred to an empty region.
    /// Calling `universal_upper_bound` for such a region gives `fr_fn_body`,
    /// which has no `external_name` in which case we use `'empty` as the
    /// region to pass to `infer_opaque_definition_from_instantiation`.
    #[instrument(skip(self, infcx))]
    pub(in crate::borrow_check) fn infer_opaque_types(
        &self,
        infcx: &InferCtxt<'_, 'tcx>,
        opaque_ty_decls: FxHashMap<DefId, ty::ResolvedOpaqueTy<'tcx>>,
        span: Span,
    ) -> FxHashMap<DefId, ty::ResolvedOpaqueTy<'tcx>> {
        opaque_ty_decls
            .into_iter()
            .map(|(opaque_def_id, ty::ResolvedOpaqueTy { concrete_type, substs })| {
                debug!(?concrete_type, ?substs);

                let mut subst_regions = vec![self.universal_regions.fr_static];
                let universal_substs =
                    infcx.tcx.fold_regions(substs, &mut false, |region, _| match *region {
                        ty::ReVar(vid) => {
                            subst_regions.push(vid);
                            self.definitions[vid].external_name.unwrap_or_else(|| {
                                infcx.tcx.sess.delay_span_bug(
                                    span,
                                    "opaque type with non-universal region substs",
                                );
                                infcx.tcx.lifetimes.re_static
                            })
                        }
                        // We don't fold regions in the predicates of opaque
                        // types to `ReVar`s. This means that in a case like
                        //
                        // fn f<'a: 'a>() -> impl Iterator<Item = impl Sized>
                        //
                        // The inner opaque type has `'static` in its substs.
                        ty::ReStatic => region,
                        _ => {
                            infcx.tcx.sess.delay_span_bug(
                                span,
                                &format!("unexpected concrete region in borrowck: {:?}", region),
                            );
                            region
                        }
                    });

                subst_regions.sort();
                subst_regions.dedup();

                let universal_concrete_type =
                    infcx.tcx.fold_regions(concrete_type, &mut false, |region, _| match *region {
                        ty::ReVar(vid) => subst_regions
                            .iter()
                            .find(|ur_vid| self.eval_equal(vid, **ur_vid))
                            .and_then(|ur_vid| self.definitions[*ur_vid].external_name)
                            .unwrap_or(infcx.tcx.lifetimes.re_root_empty),
                        ty::ReLateBound(..) => region,
                        _ => {
                            infcx.tcx.sess.delay_span_bug(
                                span,
                                &format!("unexpected concrete region in borrowck: {:?}", region),
                            );
                            region
                        }
                    });

                debug!(?universal_concrete_type, ?universal_substs);

                let remapped_type = infcx.infer_opaque_definition_from_instantiation(
                    opaque_def_id,
                    universal_substs,
                    universal_concrete_type,
                    span,
                );
                (
                    opaque_def_id,
                    ty::ResolvedOpaqueTy { concrete_type: remapped_type, substs: universal_substs },
                )
            })
            .collect()
    }

    /// Map the regions in the type to named regions. This is similar to what
    /// `infer_opaque_types` does, but can infer any universal region, not only
    /// ones from the substs for the opaque type. It also doesn't double check
    /// that the regions produced are in fact equal to the named region they are
    /// replaced with. This is fine because this function is only to improve the
    /// region names in error messages.
    pub(in crate::borrow_check) fn name_regions<T>(&self, tcx: TyCtxt<'tcx>, ty: T) -> T
    where
        T: TypeFoldable<'tcx>,
    {
        tcx.fold_regions(ty, &mut false, |region, _| match *region {
            ty::ReVar(vid) => {
                // Find something that we can name
                let upper_bound = self.approx_universal_upper_bound(vid);
                self.definitions[upper_bound].external_name.unwrap_or(region)
            }
            _ => region,
        })
    }
}
<|endoftext|>


<commit_before>// error-pattern: must be in-bounds and live at offset 5, but is outside bounds of allocation
fn main() {
    let v = [0i8; 4];
    let x = &v as *const i8;
    // The error is inside another function, so we cannot match it by line
    let x = unsafe { x.offset(5) };
    panic!("this should never print: {:?}", x);
}
<commit_msg>fix error pattern<commit_after>// error-pattern: must be in-bounds at offset 5, but is outside bounds of allocation
fn main() {
    let v = [0i8; 4];
    let x = &v as *const i8;
    // The error is inside another function, so we cannot match it by line
    let x = unsafe { x.offset(5) };
    panic!("this should never print: {:?}", x);
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 The Rouille developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

//! Apply content encodings (such as gzip compression) to the response.
//!
//! This module provides access to the content encodings supported by a request as well as
//! a function to automatically apply common content encodings to a response.
//! # Basic example
//!
//! Here is a basic example showing how to use content encodings:
//!
//! ```
//! use rouille::Request;
//! use rouille::Response;
//! use rouille::content_encoding;
//!
//! fn handle_request(request: &Request) -> Response {
//!     let response = Response::text("Hello world");
//!     content_encoding::apply(&request, response)
//! }
//! ```
use input;
use std::str;
use Request;
use Response;

/// Applies content encoding to the response.
///
/// Analyzes the `Accept-Encoding` header of the request. If one of the encodings is recognized and
/// supported by rouille, it adds a `Content-Encoding` header to the `Response` and encodes its
/// body.
///
/// If the response already has a `Content-Encoding` header, this function is a no-op.
/// If the response has a `Content-Type` header that isn't textual content, this function is a
/// no-op.
///
/// The gzip encoding is supported only if you enable the `gzip` feature of rouille (which is
/// enabled by default).
///
/// # Example
///
/// ```rust
/// use rouille::content_encoding;
/// use rouille::Request;
/// use rouille::Response;
///
/// fn handle(request: &Request) -> Response {
///     content_encoding::apply(request, Response::text("hello world"))
/// }
/// ```
pub fn apply(request: &Request, mut response: Response) -> Response {
    // Only text should be encoded. Otherwise just return.
    if !response_is_text(&response) {
        return response;
    }

    // If any of the response's headers is equal to `Content-Encoding`, ignore the function
    // call and return immediately.
    if response
        .headers
        .iter()
        .any(|&(ref key, _)| key.eq_ignore_ascii_case("Content-Encoding"))
    {
        return response;
    }

    // Now let's get the list of content encodings accepted by the request.
    // The list should be ordered from the most desired to the least desired.
    let encoding_preference = ["br", "gzip", "x-gzip", "identity"];
    let accept_encoding_header = request.header("Accept-Encoding").unwrap_or("");
    if let Some(preferred_index) = input::priority_header_preferred(
        &accept_encoding_header,
        encoding_preference.iter().cloned(),
    ) {
        match encoding_preference[preferred_index] {
            "br" => brotli(&mut response),
            "gzip" | "x-gzip" => gzip(&mut response),
            _ => (),
        }
    }

    return response;
}

// Returns true if the Content-Type of the response is a type that should be encoded.
// Since encoding is purely an optimization, it's not a problem if the function sometimes has
// false positives or false negatives.
fn response_is_text(response: &Response) -> bool {
    response.headers.iter().any(|&(ref key, ref value)| {
        if !key.eq_ignore_ascii_case("Content-Type") {
            return false;
        }

        let content_type = value.to_lowercase();
        content_type.starts_with("text/")
            || content_type.contains("javascript")
            || content_type.contains("json")
            || content_type.contains("xml")
            || content_type.contains("font")
    })
}

#[cfg(feature = "gzip")]
fn gzip(response: &mut Response) {
    use deflate::deflate_bytes_gzip;
    use std::io;
    use std::mem;
    use ResponseBody;

    response
        .headers
        .push(("Content-Encoding".into(), "gzip".into()));
    let previous_body = mem::replace(&mut response.data, ResponseBody::empty());
    let (mut raw_data, size) = previous_body.into_reader_and_size();
    let mut src = match size {
        Some(size) => Vec::with_capacity(size),
        None => Vec::new(),
    };
    io::copy(&mut raw_data, &mut src).expect("Failed reading response body while gzipping");
    let zipped = deflate_bytes_gzip(&src);
    response.data = ResponseBody::from_data(zipped);
}

#[cfg(not(feature = "gzip"))]
#[inline]
fn gzip(response: &mut Response) {}

#[cfg(feature = "brotli")]
fn brotli(response: &mut Response) {
    use brotli2::read::BrotliEncoder;
    use std::mem;
    use ResponseBody;

    response
        .headers
        .push(("Content-Encoding".into(), "br".into()));
    let previous_body = mem::replace(&mut response.data, ResponseBody::empty());
    let (raw_data, _) = previous_body.into_reader_and_size();
    response.data = ResponseBody::from_reader(BrotliEncoder::new(raw_data, 6));
}

#[cfg(not(feature = "brotli"))]
#[inline]
fn brotli(response: &mut Response) {}

#[cfg(test)]
mod tests {
    // TODO: more tests for encoding stuff
}
<commit_msg>Fix TODO: Add tests for content encoding<commit_after>// Copyright (c) 2016 The Rouille developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

//! Apply content encodings (such as gzip compression) to the response.
//!
//! This module provides access to the content encodings supported by a request as well as
//! a function to automatically apply common content encodings to a response.
//! # Basic example
//!
//! Here is a basic example showing how to use content encodings:
//!
//! ```
//! use rouille::Request;
//! use rouille::Response;
//! use rouille::content_encoding;
//!
//! fn handle_request(request: &Request) -> Response {
//!     let response = Response::text("Hello world");
//!     content_encoding::apply(&request, response)
//! }
//! ```
use input;
use Request;
use Response;

/// Applies content encoding to the response.
///
/// Analyzes the `Accept-Encoding` header of the request. If one of the encodings is recognized and
/// supported by rouille, it adds a `Content-Encoding` header to the `Response` and encodes its
/// body.
///
/// If the response already has a `Content-Encoding` header, this function is a no-op.
/// If the response has a `Content-Type` header that isn't textual content, this function is a
/// no-op.
///
/// The gzip encoding is supported only if you enable the `gzip` feature of rouille (which is
/// enabled by default).
///
/// # Example
///
/// ```rust
/// use rouille::content_encoding;
/// use rouille::Request;
/// use rouille::Response;
///
/// fn handle(request: &Request) -> Response {
///     content_encoding::apply(request, Response::text("hello world"))
/// }
/// ```
pub fn apply(request: &Request, mut response: Response) -> Response {
    // Only text should be encoded. Otherwise just return.
    if !response_is_text(&response) {
        return response;
    }

    // If any of the response's headers is equal to `Content-Encoding`, ignore the function
    // call and return immediately.
    if response
        .headers
        .iter()
        .any(|&(ref key, _)| key.eq_ignore_ascii_case("Content-Encoding"))
    {
        return response;
    }

    // Now let's get the list of content encodings accepted by the request.
    // The list should be ordered from the most desired to the least desired.
    let encoding_preference = ["br", "gzip", "x-gzip", "identity"];
    let accept_encoding_header = request.header("Accept-Encoding").unwrap_or("");
    if let Some(preferred_index) = input::priority_header_preferred(
        &accept_encoding_header,
        encoding_preference.iter().cloned(),
    ) {
        match encoding_preference[preferred_index] {
            "br" => brotli(&mut response),
            "gzip" | "x-gzip" => gzip(&mut response),
            _ => (),
        }
    }

    return response;
}

// Returns true if the Content-Type of the response is a type that should be encoded.
// Since encoding is purely an optimization, it's not a problem if the function sometimes has
// false positives or false negatives.
fn response_is_text(response: &Response) -> bool {
    response.headers.iter().any(|&(ref key, ref value)| {
        if !key.eq_ignore_ascii_case("Content-Type") {
            return false;
        }

        let content_type = value.to_lowercase();
        content_type.starts_with("text/")
            || content_type.contains("javascript")
            || content_type.contains("json")
            || content_type.contains("xml")
            || content_type.contains("font")
    })
}

#[cfg(feature = "gzip")]
fn gzip(response: &mut Response) {
    use deflate::deflate_bytes_gzip;
    use std::io;
    use std::mem;
    use ResponseBody;

    response
        .headers
        .push(("Content-Encoding".into(), "gzip".into()));
    let previous_body = mem::replace(&mut response.data, ResponseBody::empty());
    let (mut raw_data, size) = previous_body.into_reader_and_size();
    let mut src = match size {
        Some(size) => Vec::with_capacity(size),
        None => Vec::new(),
    };
    io::copy(&mut raw_data, &mut src).expect("Failed reading response body while gzipping");
    let zipped = deflate_bytes_gzip(&src);
    response.data = ResponseBody::from_data(zipped);
}

#[cfg(not(feature = "gzip"))]
#[inline]
fn gzip(response: &mut Response) {}

#[cfg(feature = "brotli")]
fn brotli(response: &mut Response) {
    use brotli2::read::BrotliEncoder;
    use std::mem;
    use ResponseBody;

    response
        .headers
        .push(("Content-Encoding".into(), "br".into()));
    let previous_body = mem::replace(&mut response.data, ResponseBody::empty());
    let (raw_data, _) = previous_body.into_reader_and_size();
    response.data = ResponseBody::from_reader(BrotliEncoder::new(raw_data, 6));
}

#[cfg(not(feature = "brotli"))]
#[inline]
fn brotli(response: &mut Response) {}

#[cfg(test)]
mod tests {
    use content_encoding;
    use Request;
    use Response;

    // TODO: more tests for encoding stuff
    #[test]
    fn text_response() {
        assert!(content_encoding::response_is_text(&Response::text("")));
    }

    #[test]
    fn non_text_response() {
        assert!(!content_encoding::response_is_text(&Response::from_data(
            "image/jpeg",
            ""
        )));
    }

    #[test]
    fn no_req_encodings() {
        let request = Request::fake_http("GET", "/", vec![], vec![]);
        let response = Response::html("
Hello world
");
        let encoded_response = content_encoding::apply(&request, response);
        assert!(!encoded_response
            .headers
            .iter()
            .any(|(header_name, _)| header_name == "Content-Encoding")); // No Content-Encoding header
        let mut encoded_content = vec![];
        encoded_response
            .data
            .into_reader_and_size()
            .0
            .read_to_end(&mut encoded_content)
            .unwrap();
        assert_eq!(
            String::from_utf8(encoded_content).unwrap(),
            "
Hello world
"
        ); // No encoding applied
    }

    #[test]
    fn empty_req_encodings() {
        let request = {
            let h = vec![("Accept-Encoding".to_owned(), "".to_owned())];
            Request::fake_http("GET", "/", h, vec![])
        };
        let response = Response::html("
Hello world
");

        let encoded_response = content_encoding::apply(&request, response);
        assert!(!encoded_response
            .headers
            .iter()
            .any(|(header_name, _)| header_name == "Content-Encoding")); // No Content-Encoding header
        let mut encoded_content = vec![];
        encoded_response
            .data
            .into_reader_and_size()
            .0
            .read_to_end(&mut encoded_content)
            .unwrap();
        assert_eq!(
            String::from_utf8(encoded_content).unwrap(),
            "
Hello world
"
        ); // No encoding applied
    }

    #[test]
    fn multi_req_encoding() {
        let request = {
            let h = vec![("Accept-Encoding".to_owned(), "foo".to_owned())];
            Request::fake_http("GET", "/", h, vec![])
        };
        let response = Response::html("
Hello world
");

        let encoded_response = content_encoding::apply(&request, response);
        assert!(!encoded_response
            .headers
            .iter()
            .any(|(header_name, _)| header_name == "Content-Encoding")); // No Content-Encoding header
        let mut encoded_content = vec![];
        encoded_response
            .data
            .into_reader_and_size()
            .0
            .read_to_end(&mut encoded_content)
            .unwrap();
        assert_eq!(
            String::from_utf8(encoded_content).unwrap(),
            "
Hello world
"
        ); // No encoding applied
    }

    #[test]
    fn unknown_req_encoding() {
        let request = {
            let h = vec![("Accept-Encoding".to_owned(), "x-gzip, br".to_owned())];
            Request::fake_http("GET", "/", h, vec![])
        };
        let response = Response::html("
Hello world
");

        let encoded_response = content_encoding::apply(&request, response);
        assert!(encoded_response
            .headers
            .contains(&("Content-Encoding".into(), "br".into()))); // Brotli Content-Encoding header
    }

    #[test]
    fn brotli_encoding() {
        let request = {
            let h = vec![("Accept-Encoding".to_owned(), "br".to_owned())];
            Request::fake_http("GET", "/", h, vec![])
        };
        let response = Response::html(
            "<title>Hello world</title>
Hello world
",
        );

        let encoded_response = content_encoding::apply(&request, response);
        assert!(encoded_response
            .headers
            .contains(&("Content-Encoding".into(), "br".into()))); // Brotli Content-Encoding header
        let mut encoded_content = vec![];
        encoded_response
            .data
            .into_reader_and_size()
            .0
            .read_to_end(&mut encoded_content)
            .unwrap();
        assert_eq!(
            encoded_content,
            vec![
                27, 75, 0, 0, 4, 28, 114, 164, 129, 5, 210, 206, 25, 30, 90, 114, 224, 114, 73,
                109, 45, 196, 23, 126, 240, 144, 77, 40, 26, 211, 228, 67, 73, 40, 236, 55, 101,
                254, 127, 147, 194, 129, 132, 65, 130, 120, 152, 249, 68, 56, 93, 2
            ]
        ); // Applied proper Brotli encoding
    }

    #[test]
    fn gzip_encoding() {
        let request = {
            let h = vec![("Accept-Encoding".to_owned(), "gzip".to_owned())];
            Request::fake_http("GET", "/", h, vec![])
        };
        let response = Response::html(
            "<title>Hello world</title>
Hello world
",
        );

        let encoded_response = content_encoding::apply(&request, response);
        assert!(encoded_response
            .headers
            .contains(&("Content-Encoding".into(), "gzip".into()))); // gzip Content-Encoding header
        let mut encoded_content = vec![];
        encoded_response
            .data
            .into_reader_and_size()
            .0
            .read_to_end(&mut encoded_content)
            .unwrap();
        assert_eq!(
            encoded_content,
            vec![
                31, 139, 8, 0, 0, 0, 0, 0, 0, 3, 179, 201, 40, 201, 205, 177, 179, 201, 72, 77, 76,
                177, 179, 41, 201, 44, 201, 73, 181, 243, 72, 205, 201, 201, 87, 40, 207, 47, 202,
                73, 177, 209, 135, 8, 217, 36, 229, 167, 84, 218, 217, 20, 160, 202, 21, 216, 217,
                232, 67, 36, 244, 193, 166, 0, 0, 202, 239, 44, 120, 76, 0, 0, 0
            ]
        ); // Applied proper gzip encoding
    }
}
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

//! Rosetta Account API
//!
//! See: [Account API Spec](https://www.rosetta-api.org/docs/AccountApi.html)
//!

use crate::types::*;
use crate::{
    common::{
        check_network, get_block_index_from_request, handle_request, native_coin, native_coin_tag,
        with_context,
    },
    error::{ApiError, ApiResult},
    types::{AccountBalanceRequest, AccountBalanceResponse, Amount, Currency, CurrencyMetadata},
    RosettaContext,
};
use aptos_logger::{debug, trace};
use aptos_sdk::move_types::language_storage::TypeTag;
use aptos_types::account_address::AccountAddress;
use aptos_types::account_config::{AccountResource, CoinInfoResource, CoinStoreResource};
use aptos_types::stake_pool::StakePool;
use once_cell::sync::Lazy;
use percent_encoding::{utf8_percent_encode, AsciiSet, CONTROLS};
use std::collections::BTreeMap;
use std::{
    collections::HashSet,
    sync::{Arc, RwLock},
};
use warp::Filter;

/// Account routes e.g. balance
pub fn routes(
    server_context: RosettaContext,
) -> impl Filter<Extract = impl warp::Reply, Error = warp::Rejection> + Clone {
    warp::post().and(
        warp::path!("account" / "balance")
            .and(warp::body::json())
            .and(with_context(server_context))
            .and_then(handle_request(account_balance)),
    )
}

/// Account balance command
///
/// [API Spec](https://www.rosetta-api.org/docs/AccountApi.html#accountbalance)
async fn account_balance(
    request: AccountBalanceRequest,
    server_context: RosettaContext,
) -> ApiResult<AccountBalanceResponse> {
    debug!("/account/balance");
    trace!(
        request = ?request,
        server_context = ?server_context,
        "account_balance for [{}]",
        request.account_identifier.address
    );

    let network_identifier = request.network_identifier;

    check_network(network_identifier, &server_context)?;
    let rest_client = server_context.rest_client()?;

    // Retrieve the block index to read
    let block_height =
        get_block_index_from_request(&server_context, request.block_identifier.clone()).await?;

    // Version to grab is the last entry in the block (balance is at end of block)
    let block_info = server_context
        .block_cache()?
        .get_block_info_by_height(block_height, server_context.chain_id)
        .await?;
    let balance_version = block_info.last_version;

    let (sequence_number, balances) = get_balances(
        server_context.coin_cache.clone(),
        &rest_client,
        request.account_identifier,
        balance_version,
        request.currencies,
    )
    .await?;

    Ok(AccountBalanceResponse {
        block_identifier: block_info.block_id,
        balances,
        metadata: AccountBalanceMetadata {
            sequence_number: sequence_number.into(),
        },
    })
}

/// Retrieve the balances for an account
async fn get_balances(
    coin_cache: Arc<CoinCache>,
    rest_client: &aptos_rest_client::Client,
    account: AccountIdentifier,
    version: u64,
    maybe_filter_currencies: Option<Vec<Currency>>,
) -> ApiResult<(u64, Vec<Amount>)> {
    let owner_address = account.account_address()?;

    // Retrieve all account resources
    if let Ok(response) = rest_client
        .get_account_resources_at_version_bcs(owner_address, version)
        .await
    {
        let resources = response.into_inner();
        let mut maybe_sequence_number = None;
        let mut balances = vec![];

        // Iterate through resources, converting balances
        for (struct_tag, bytes) in resources {
            match (
                struct_tag.address,
                struct_tag.module.as_str(),
                struct_tag.name.as_str(),
            ) {
                (AccountAddress::ONE, ACCOUNT_MODULE, ACCOUNT_RESOURCE) => {
                    let account: AccountResource = bcs::from_bytes(&bytes)?;
                    maybe_sequence_number = Some(account.sequence_number())
                }
                (AccountAddress::ONE, COIN_MODULE, COIN_STORE_RESOURCE) => {
                    // Only show coins on the base account
                    if account.is_base_account() {
                        let coin_store: CoinStoreResource = bcs::from_bytes(&bytes)?;
                        if let Some(coin_type) = struct_tag.type_params.first() {
                            // Only display supported coins
                            if let Some(currency) =
                                coin_cache.get_currency(coin_type, Some(version)).await?
                            {
                                balances.push(Amount {
                                    value: coin_store.coin().to_string(),
                                    currency,
                                });
                            }
                        }
                    }
                }
                (AccountAddress::ONE, STAKE_MODULE, STAKE_POOL_RESOURCE) => {
                    if account.is_base_account() {
                        continue;
                    }

                    let stake_pool: StakePool = bcs::from_bytes(&bytes)?;
                    // If the operator address is different, skip
                    if account.is_operator_stake()
                        && account.operator_address()? != stake_pool.operator_address
                    {
                        continue;
                    }

                    // TODO: Represent inactive, and pending as separate?
                    let balance = stake_pool.active
                        + stake_pool.inactive
                        + stake_pool.pending_active
                        + stake_pool.pending_inactive;
                    balances.push(Amount {
                        value: balance.to_string(),
                        currency: native_coin(),
                    })
                }
                (AccountAddress::ONE, STAKING_CONTRACT_MODULE, STORE_RESOURCE) => {
                    if account.is_base_account() {
                        continue;
                    }

                    let store: Store = bcs::from_bytes(&bytes)?;
                    if account.is_total_stake() {
                        // For total stake, collect all underlying staking contracts and combine
                        let mut total_stake: Option<u64> = None;
                        for (_operator, contract) in store.staking_contracts {
                            total_stake =
                                Some(total_stake.unwrap_or_default() + contract.principal);
                        }

                        if let Some(balance) = total_stake {
                            balances.push(Amount {
                                value: balance.to_string(),
                                currency: native_coin(),
                            })
                        }
                    } else if account.is_operator_stake() {
                        // For operator stake, filter on operator address
                        let operator_address = account.operator_address()?;
                        if let Some(contract) = store.staking_contracts.get(&operator_address) {
                            balances.push(Amount {
                                value: contract.principal.to_string(),
                                currency: native_coin(),
                            })
                        }
                    }
                }
                _ => {}
            }
        }

        let sequence_number = if let Some(sequence_number) = maybe_sequence_number {
            sequence_number
        } else {
            return Err(ApiError::InternalError(Some(
                "Failed to retrieve account sequence number".to_string(),
            )));
        };

        // Filter based on requested currencies
        if let Some(currencies) = maybe_filter_currencies {
            let mut currencies: HashSet<Currency> = currencies.into_iter().collect();
            // Remove extra currencies not requested
            balances = balances
                .into_iter()
                .filter(|balance| currencies.contains(&balance.currency))
                .collect();

            for balance in balances.iter() {
                currencies.remove(&balance.currency);
            }

            for currency in currencies {
                balances.push(Amount {
                    value: 0.to_string(),
                    currency,
                });
            }
        }

        // Retrieve balances
        Ok((sequence_number, balances))
    } else {
        Ok((
            0,
            vec![Amount {
                value: 0.to_string(),
                currency: native_coin(),
            }],
        ))
    }
}

/// Only coins supported by Rosetta
///
/// TODO: Allow passing in known supported coins in some sort of config file
pub static SUPPORTED_COINS: Lazy<Vec<TypeTag>> = Lazy::new(|| vec![native_coin_tag()]);

/// A cache for currencies, so we don't have to keep looking up the status of it
#[derive(Debug)]
pub struct CoinCache {
    rest_client: Option<Arc<aptos_rest_client::Client>>,
    currencies: RwLock<BTreeMap<TypeTag, Option<Currency>>>,
}

impl CoinCache {
    pub fn new(rest_client: Option<Arc<aptos_rest_client::Client>>) -> Self {
        Self {
            rest_client,
            currencies: RwLock::new(BTreeMap::new()),
        }
    }

    pub fn get_currency_from_cache(&self, coin: &TypeTag) -> Option<Currency> {
        // Short circuit for the default coin
        if coin == &native_coin_tag() {
            return Some(native_coin());
        }

        let currencies = self
            .currencies
            .read()
            .expect("Can't recover from poisoned lock on coin cache");

        if let Some(currency) = currencies.get(coin) {
            currency.clone()
        } else {
            None
        }
    }

    /// Retrieve a currency and cache it if applicable
    pub async fn get_currency(
        &self,
        coin: &TypeTag,
        version: Option<u64>,
    ) -> ApiResult<Option<Currency>> {
        // Short circuit for the default coin
        if coin == &native_coin_tag() {
            return Ok(Some(native_coin()));
        }

        // Unsupported coins should be ignored, as symbol may be used as an identifier
        if !SUPPORTED_COINS.contains(coin) {
            return Ok(None);
        }

        {
            let currencies = self
                .currencies
                .read()
                .expect("Can't recover from poisoned lock on coin cache");
            if let Some(currency) = currencies.get(coin) {
                return Ok(currency.clone());
            }
        }

        let currency = self
            .get_currency_inner(coin, version)
            .await
            .map_err(|err| {
                ApiError::CoinTypeFailedToBeFetched(Some(format!(
                    "Failed to retrieve coin type {} {}",
                    coin, err
                )))
            })?;
        self.currencies
            .write()
            .expect("Can't recover from poisoned lock on coin cache")
            .insert(coin.clone(), currency.clone());
        Ok(currency)
    }

    /// Pulls currency information from onchain
    pub async fn get_currency_inner(
        &self,
        coin: &TypeTag,
        version: Option<u64>,
    ) -> ApiResult<Option<Currency>> {
        let struct_tag = match coin {
            TypeTag::Struct(ref tag) => tag,
            // This is a poorly formed coin, and we'll just skip over it
            _ => return Ok(None),
        };

        // Retrieve the coin type
        const ENCODE_CHARS: &AsciiSet = &CONTROLS.add(b'<').add(b'>');
        let address = struct_tag.address;
        let resource_tag = format!("0x1::coin::CoinInfo<{}>", struct_tag);
        let encoded_resource_tag = utf8_percent_encode(&resource_tag, ENCODE_CHARS).to_string();

        let response = if let Some(version) = version {
            self.rest_client
                .as_ref()
                .expect("Should have rest client to check coin cache")
                .get_account_resource_at_version_bcs::<CoinInfoResource>(
                    address,
                    &encoded_resource_tag,
                    version,
                )
                .await
        } else {
            self.rest_client
                .as_ref()
                .expect("Should have rest client to check coin cache")
                .get_account_resource_bcs::<CoinInfoResource>(address, &encoded_resource_tag)
                .await
        };

        // At this point if we've retrieved it and it's bad, we error out
        let coin_info = response?.into_inner();

        // TODO: The symbol has to come from a trusted list, as the move type is the real indicator
        Ok(Some(Currency {
            symbol: coin_info
                .symbol()
                .unwrap_or_else(|err| format!("Invalid Symbol: {}", err)),
            decimals: coin_info.decimals(),
            metadata: Some(CurrencyMetadata {
                move_type: struct_tag.to_string(),
            }),
        }))
    }
}
<commit_msg>[rosetta] Use stake amounts rather than principal for stake pool balances<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

//! Rosetta Account API
//!
//! See: [Account API Spec](https://www.rosetta-api.org/docs/AccountApi.html)
//!

use crate::types::*;
use crate::{
    common::{
        check_network, get_block_index_from_request, handle_request, native_coin, native_coin_tag,
        with_context,
    },
    error::{ApiError, ApiResult},
    types::{AccountBalanceRequest, AccountBalanceResponse, Amount, Currency, CurrencyMetadata},
    RosettaContext,
};
use aptos_logger::{debug, trace};
use aptos_sdk::move_types::language_storage::TypeTag;
use aptos_types::account_address::AccountAddress;
use aptos_types::account_config::{AccountResource, CoinInfoResource, CoinStoreResource};
use aptos_types::stake_pool::StakePool;
use once_cell::sync::Lazy;
use percent_encoding::{utf8_percent_encode, AsciiSet, CONTROLS};
use std::collections::BTreeMap;
use std::str::FromStr;
use std::{
    collections::HashSet,
    sync::{Arc, RwLock},
};
use warp::Filter;

/// Account routes e.g. balance
pub fn routes(
    server_context: RosettaContext,
) -> impl Filter<Extract = impl warp::Reply, Error = warp::Rejection> + Clone {
    warp::post().and(
        warp::path!("account" / "balance")
            .and(warp::body::json())
            .and(with_context(server_context))
            .and_then(handle_request(account_balance)),
    )
}

/// Account balance command
///
/// [API Spec](https://www.rosetta-api.org/docs/AccountApi.html#accountbalance)
async fn account_balance(
    request: AccountBalanceRequest,
    server_context: RosettaContext,
) -> ApiResult<AccountBalanceResponse> {
    debug!("/account/balance");
    trace!(
        request = ?request,
        server_context = ?server_context,
        "account_balance for [{}]",
        request.account_identifier.address
    );

    let network_identifier = request.network_identifier;

    check_network(network_identifier, &server_context)?;
    let rest_client = server_context.rest_client()?;

    // Retrieve the block index to read
    let block_height =
        get_block_index_from_request(&server_context, request.block_identifier.clone()).await?;

    // Version to grab is the last entry in the block (balance is at end of block)
    let block_info = server_context
        .block_cache()?
        .get_block_info_by_height(block_height, server_context.chain_id)
        .await?;
    let balance_version = block_info.last_version;

    let (sequence_number, balances) = get_balances(
        server_context.coin_cache.clone(),
        &rest_client,
        request.account_identifier,
        balance_version,
        request.currencies,
    )
    .await?;

    Ok(AccountBalanceResponse {
        block_identifier: block_info.block_id,
        balances,
        metadata: AccountBalanceMetadata {
            sequence_number: sequence_number.into(),
        },
    })
}

/// Retrieve the balances for an account
async fn get_balances(
    coin_cache: Arc<CoinCache>,
    rest_client: &aptos_rest_client::Client,
    account: AccountIdentifier,
    version: u64,
    maybe_filter_currencies: Option<Vec<Currency>>,
) -> ApiResult<(u64, Vec<Amount>)> {
    let owner_address = account.account_address()?;

    // Retrieve all account resources
    if let Ok(response) = rest_client
        .get_account_resources_at_version_bcs(owner_address, version)
        .await
    {
        let resources = response.into_inner();
        let mut maybe_sequence_number = None;
        let mut balances = vec![];

        // Iterate through resources, converting balances
        for (struct_tag, bytes) in resources {
            match (
                struct_tag.address,
                struct_tag.module.as_str(),
                struct_tag.name.as_str(),
            ) {
                (AccountAddress::ONE, ACCOUNT_MODULE, ACCOUNT_RESOURCE) => {
                    let account: AccountResource = bcs::from_bytes(&bytes)?;
                    maybe_sequence_number = Some(account.sequence_number())
                }
                (AccountAddress::ONE, COIN_MODULE, COIN_STORE_RESOURCE) => {
                    // Only show coins on the base account
                    if account.is_base_account() {
                        let coin_store: CoinStoreResource = bcs::from_bytes(&bytes)?;
                        if let Some(coin_type) = struct_tag.type_params.first() {
                            // Only display supported coins
                            if let Some(currency) =
                                coin_cache.get_currency(coin_type, Some(version)).await?
                            {
                                balances.push(Amount {
                                    value: coin_store.coin().to_string(),
                                    currency,
                                });
                            }
                        }
                    }
                }
                (AccountAddress::ONE, STAKE_MODULE, STAKE_POOL_RESOURCE) => {
                    let stake_pool: StakePool = bcs::from_bytes(&bytes)?;
                    if let Ok(balance) = get_stake_balance_from_stake_pool(&stake_pool, &account) {
                        balances.push(balance);
                    }
                }
                (AccountAddress::ONE, STAKING_CONTRACT_MODULE, STORE_RESOURCE) => {
                    if account.is_base_account() {
                        continue;
                    }

                    let store: Store = bcs::from_bytes(&bytes)?;
                    if account.is_total_stake() {
                        // For total stake, collect all underlying staking contracts and combine
                        let mut total_stake: Option<u64> = None;
                        for (_operator, contract) in store.staking_contracts {
                            if let Ok(response) = rest_client
                                .get_account_resource_bcs::<StakePool>(
                                    contract.pool_address,
                                    "0x1::stake::StakePool",
                                )
                                .await
                            {
                                let stake_pool = response.into_inner();

                                // Any stake pools that match, retrieve that.  Then update the total
                                if let Ok(balance) =
                                    get_stake_balance_from_stake_pool(&stake_pool, &account)
                                {
                                    total_stake = Some(
                                        total_stake.unwrap_or_default()
                                            + u64::from_str(&balance.value).unwrap_or_default(),
                                    );
                                }
                            }
                        }

                        if let Some(balance) = total_stake {
                            balances.push(Amount {
                                value: balance.to_string(),
                                currency: native_coin(),
                            })
                        }
                    } else if account.is_operator_stake() {
                        // For operator stake, filter on operator address
                        let operator_address = account.operator_address()?;
                        if let Some(contract) = store.staking_contracts.get(&operator_address) {
                            if let Ok(response) = rest_client
                                .get_account_resource_bcs::<StakePool>(
                                    contract.pool_address,
                                    "0x1::stake::StakePool",
                                )
                                .await
                            {
                                let stake_pool = response.into_inner();

                                // Any stake pools that match, retrieve that.  Then update the total
                                if let Ok(balance) =
                                    get_stake_balance_from_stake_pool(&stake_pool, &account)
                                {
                                    balances.push(balance)
                                }
                            }
                        }
                    }
                }
                _ => {}
            }
        }

        let sequence_number = if let Some(sequence_number) = maybe_sequence_number {
            sequence_number
        } else {
            return Err(ApiError::InternalError(Some(
                "Failed to retrieve account sequence number".to_string(),
            )));
        };

        // Filter based on requested currencies
        if let Some(currencies) = maybe_filter_currencies {
            let mut currencies: HashSet<Currency> = currencies.into_iter().collect();
            // Remove extra currencies not requested
            balances = balances
                .into_iter()
                .filter(|balance| currencies.contains(&balance.currency))
                .collect();

            for balance in balances.iter() {
                currencies.remove(&balance.currency);
            }

            for currency in currencies {
                balances.push(Amount {
                    value: 0.to_string(),
                    currency,
                });
            }
        }

        // Retrieve balances
        Ok((sequence_number, balances))
    } else {
        Ok((
            0,
            vec![Amount {
                value: 0.to_string(),
                currency: native_coin(),
            }],
        ))
    }
}

/// Retrieves total stake balances from an individual stake pool
fn get_stake_balance_from_stake_pool(
    stake_pool: &StakePool,
    account: &AccountIdentifier,
) -> ApiResult<Amount> {
    // Stake isn't allowed for base accounts
    if account.is_base_account() {
        return Err(ApiError::InvalidInput(Some(
            "Stake pool not supported for base account".to_string(),
        )));
    }

    // If the operator address is different, skip
    if account.is_operator_stake() && account.operator_address()? != stake_pool.operator_address {
        return Err(ApiError::InvalidInput(Some(
            "Stake pool not for matching operator".to_string(),
        )));
    }

    // TODO: Represent inactive, and pending as separate?
    let balance = stake_pool.active
        + stake_pool.inactive
        + stake_pool.pending_active
        + stake_pool.pending_inactive;
    Ok(Amount {
        value: balance.to_string(),
        currency: native_coin(),
    })
}

/// Only coins supported by Rosetta
///
/// TODO: Allow passing in known supported coins in some sort of config file
pub static SUPPORTED_COINS: Lazy<Vec<TypeTag>> = Lazy::new(|| vec![native_coin_tag()]);

/// A cache for currencies, so we don't have to keep looking up the status of it
#[derive(Debug)]
pub struct CoinCache {
    rest_client: Option<Arc<aptos_rest_client::Client>>,
    currencies: RwLock<BTreeMap<TypeTag, Option<Currency>>>,
}

impl CoinCache {
    pub fn new(rest_client: Option<Arc<aptos_rest_client::Client>>) -> Self {
        Self {
            rest_client,
            currencies: RwLock::new(BTreeMap::new()),
        }
    }

    pub fn get_currency_from_cache(&self, coin: &TypeTag) -> Option<Currency> {
        // Short circuit for the default coin
        if coin == &native_coin_tag() {
            return Some(native_coin());
        }

        let currencies = self
            .currencies
            .read()
            .expect("Can't recover from poisoned lock on coin cache");

        if let Some(currency) = currencies.get(coin) {
            currency.clone()
        } else {
            None
        }
    }

    /// Retrieve a currency and cache it if applicable
    pub async fn get_currency(
        &self,
        coin: &TypeTag,
        version: Option<u64>,
    ) -> ApiResult<Option<Currency>> {
        // Short circuit for the default coin
        if coin == &native_coin_tag() {
            return Ok(Some(native_coin()));
        }

        // Unsupported coins should be ignored, as symbol may be used as an identifier
        if !SUPPORTED_COINS.contains(coin) {
            return Ok(None);
        }

        {
            let currencies = self
                .currencies
                .read()
                .expect("Can't recover from poisoned lock on coin cache");
            if let Some(currency) = currencies.get(coin) {
                return Ok(currency.clone());
            }
        }

        let currency = self
            .get_currency_inner(coin, version)
            .await
            .map_err(|err| {
                ApiError::CoinTypeFailedToBeFetched(Some(format!(
                    "Failed to retrieve coin type {} {}",
                    coin, err
                )))
            })?;
        self.currencies
            .write()
            .expect("Can't recover from poisoned lock on coin cache")
            .insert(coin.clone(), currency.clone());
        Ok(currency)
    }

    /// Pulls currency information from onchain
    pub async fn get_currency_inner(
        &self,
        coin: &TypeTag,
        version: Option<u64>,
    ) -> ApiResult<Option<Currency>> {
        let struct_tag = match coin {
            TypeTag::Struct(ref tag) => tag,
            // This is a poorly formed coin, and we'll just skip over it
            _ => return Ok(None),
        };

        // Retrieve the coin type
        const ENCODE_CHARS: &AsciiSet = &CONTROLS.add(b'<').add(b'>');
        let address = struct_tag.address;
        let resource_tag = format!("0x1::coin::CoinInfo<{}>", struct_tag);
        let encoded_resource_tag = utf8_percent_encode(&resource_tag, ENCODE_CHARS).to_string();

        let response = if let Some(version) = version {
            self.rest_client
                .as_ref()
                .expect("Should have rest client to check coin cache")
                .get_account_resource_at_version_bcs::<CoinInfoResource>(
                    address,
                    &encoded_resource_tag,
                    version,
                )
                .await
        } else {
            self.rest_client
                .as_ref()
                .expect("Should have rest client to check coin cache")
                .get_account_resource_bcs::<CoinInfoResource>(address, &encoded_resource_tag)
                .await
        };

        // At this point if we've retrieved it and it's bad, we error out
        let coin_info = response?.into_inner();

        // TODO: The symbol has to come from a trusted list, as the move type is the real indicator
        Ok(Some(Currency {
            symbol: coin_info
                .symbol()
                .unwrap_or_else(|err| format!("Invalid Symbol: {}", err)),
            decimals: coin_info.decimals(),
            metadata: Some(CurrencyMetadata {
                move_type: struct_tag.to_string(),
            }),
        }))
    }
}
<|endoftext|>