<commit_before>// This function is unstable yet, tests ignored

use crate::index::btree::node::read_unchecked;
use crate::index::btree::node::write_node;
use crate::index::btree::node::write_targeted;
use crate::index::btree::node::EmptyNode;
use crate::index::btree::node::NodeData;
use crate::index::btree::node::NodeWriteGuard;
use crate::index::btree::search::mut_search;
use crate::index::btree::search::MutSearchResult;
use crate::index::btree::NodeCellRef;
use crate::index::btree::{external, BPlusTree};
use crate::index::trees::EntryKey;
use crate::index::trees::Slice;
use std::fmt::Debug;
use std::mem;

pub struct RebalancingNodes<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub left_guard: NodeWriteGuard<KS, PS>,
    pub left_ref: NodeCellRef,
    pub right_guard: NodeWriteGuard<KS, PS>,
    pub right_right_guard: NodeWriteGuard<KS, PS>, // for external pointer modification
    pub parent: NodeWriteGuard<KS, PS>,
    pub parent_pos: usize,
}

pub struct RemoveResult<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub rebalancing: Option<RebalancingNodes<KS, PS>>,
    pub removed: bool,
}

pub fn with_innode_removing<F, KS, PS>(
    key: &EntryKey,
    mut rebalancing: RebalancingNodes<KS, PS>,
    parent: &NodeCellRef,
    parent_parent: &NodeCellRef,
    removed: bool,
    level: usize,
    func: F,
) -> RemoveResult<KS, PS>
where
    F: Fn(&mut RebalancingNodes<KS, PS>),
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let parent_half_full = rebalancing.parent.is_half_full();
    let mut parent_right_guard = write_node(&rebalancing.parent.innode_mut().right);
    let mut parent_remove_result = RemoveResult {
        rebalancing: None,
        removed,
    };

    let parent_parent_guard;
    let parent_parent_remove_pos;
    let parent_right_right_guard;
    if level == 0 {
        parent_parent_guard = write_node(parent_parent);
        parent_right_right_guard = None;
        parent_parent_remove_pos = 0;
    } else {
        let pre_locked_parent_right_right_guard = parent_right_guard
            .right_ref_mut_no_empty()
            .map(|r| write_node(r));
        let pp_guard = write_targeted(write_node(parent_parent), key);
        parent_parent_guard = pp_guard;
        parent_parent_remove_pos = parent_parent_guard.search(key);
        let parent_right_half_full = parent_right_guard.is_half_full();
        parent_right_right_guard = if parent_parent_guard.len() > parent_parent_remove_pos + 1
            && (!parent_half_full || !parent_right_half_full)
        {
            // indicates whether the upper parent level need to relocated
            pre_locked_parent_right_right_guard
        } else {
            None
        };
    }

    func(&mut rebalancing);
    parent_remove_result.rebalancing =
        parent_right_right_guard.map(|parent_right_right_guard_stripped| RebalancingNodes {
            left_guard: rebalancing.parent,
            left_ref: parent.clone(),
            right_right_guard: parent_right_right_guard_stripped,
            right_guard: parent_right_guard,
            parent: parent_parent_guard,
            parent_pos: parent_parent_remove_pos,
        });
    parent_remove_result
}

pub fn remove_from_node<KS, PS>(
    tree: &BPlusTree<KS, PS>,
    node_ref: &mut NodeCellRef,
    key: &mut EntryKey,
    parent: &NodeCellRef,
    level: usize,
) -> RemoveResult<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    debug!("Removing {:?} from node, level {}", key, level);
    let search = mut_search::<KS, PS>(node_ref, key);
    match search {
        MutSearchResult::Internal(mut sub_node) => {
            let node_remove_res =
                remove_from_node(tree, &mut sub_node, key, node_ref, level + 1);
            let removed = node_remove_res.removed;
            if let Some(mut rebalancing) = node_remove_res.rebalancing {
                debug!(
                    "Need to rebalance sub nodes, key {:?}, level {}",
                    key, level
                );
                if rebalancing.left_guard.is_empty() && !rebalancing.right_guard.is_empty() {
                    debug!(
                        "Remove {:?} sub level left node is empty, level {}",
                        key, level
                    );
                    Some(with_innode_removing(
                        key,
                        rebalancing,
                        &sub_node,
                        parent,
                        removed,
                        level,
                        |rebalancing| {
                            // Remove the empty node that have a right node with the same parent
                            // Because we cannot lock from left to right, we have to move the content
                            // of the right node to the left and remove the right node instead so the
                            // left right pointers can be modified
                            if !rebalancing.left_guard.is_empty()
                                || (level == 0
                                    && read_unchecked::<KS, PS>(&(*tree.root.read())).first_key()
                                        != rebalancing.left_guard.first_key())
                            {
                                return;
                            }

                            let left_node = &mut *rebalancing.left_guard;
                            let right_node = &mut *rebalancing.right_guard;
                            // swap the content of left and right, then delete right
                            // this procedure will prevent locking left node for changing right reference
                            let left_left_ref = left_node.left_ref_mut().map(|r| r.clone());
                            let right_right_ref =
                                right_node.right_ref_mut_no_empty().map(|r| r.clone());
                            let left_ref = rebalancing.left_ref.clone();
                            // swap the empty node with the right node. In this case left node holds
                            // content of the right node but pointers need to be corrected.
                            mem::swap(left_node, right_node);
                            *right_node = NodeData::Empty(box EmptyNode {
                                left: Some(left_ref.clone()),
                                right: left_ref.clone(),
                            });
                            left_node
                                .left_ref_mut()
                                .map(|r| *r = left_left_ref.unwrap());
                            left_node
                                .right_ref_mut_no_empty()
                                .map(|r| *r = right_right_ref.unwrap());
                            rebalancing
                                .right_right_guard
                                .left_ref_mut()
                                .map(|r| *r = left_ref);
                            // remove the left ptr
                            rebalancing.parent.remove(rebalancing.parent_pos);
                            // point the right ptr to the replaced sub node
                            debug!(
                                "Removed sub level empty node, living node len {}, have {:?}",
                                left_node.len(),
                                left_node.keys()
                            );
                            rebalancing.parent.innode_mut().ptrs.as_slice()
                                [rebalancing.parent_pos + 1] = sub_node.clone();
                        },
                    ))
                } else if rebalancing.right_guard.is_empty() {
                    debug!(
                        "Remove {:?} sub level right node is empty, level {}",
                        key, level
                    );
                    Some(with_innode_removing(
                        key,
                        rebalancing,
                        &sub_node,
                        parent,
                        removed,
                        level,
                        |rebalancing| {
                            // There is a right empty node that can be deleted directly without hassle
                            let node_to_remove = &mut rebalancing.right_guard;
                            let owned_left_ref = rebalancing.left_ref.clone();
                            rebalancing.left_guard.right_ref_mut_no_empty().map(|r| {
                                *r = node_to_remove.right_ref_mut_no_empty().unwrap().clone()
                            });
                            rebalancing
                                .right_right_guard
                                .left_ref_mut()
                                .map(|r| *r = owned_left_ref);
                            rebalancing.parent.remove(rebalancing.parent_pos + 1);
                        },
                    ))
                } else if rebalancing.left_guard.cannot_merge()
                    || rebalancing.right_guard.cannot_merge()
                {
                    if rebalancing.right_guard.len() < KS::slice_len() / 3
                        && rebalancing.right_guard.len() < rebalancing.left_guard.len()
                        && rebalancing.left_guard.len() > KS::slice_len() / 2
                    {
                        // Relocate the nodes with the same parent for balance.
                        // For OLFIT, items can only be located from left to right.
                        debug!(
                            "Remove {:?} sub level need relocation, level {}",
                            key, level
                        );
                        let left_node = &mut rebalancing.left_guard;
                        let right_node = &mut rebalancing.right_guard;
                        let left_pos = rebalancing.parent_pos;
                        let right_pos = left_pos + 1;
                        rebalancing
                            .parent
                            .innode_mut()
                            .relocate_children(left_pos, right_pos, left_node, right_node, tree);
                    }
                    None
                } else if rebalancing.left_guard.len() + rebalancing.right_guard.len() + 1
                    <= KS::slice_len()
                {
                    // Nodes with the same parent can merge
                    debug!(
                        "Remove {:?} sub level need to be merged, level {}",
                        key, level
                    );
                    Some(with_innode_removing(
                        key,
                        rebalancing,
                        &sub_node,
                        parent,
                        removed,
                        level,
                        |rebalancing| {
                            let left_node = &mut rebalancing.left_guard;
                            let right_node = &mut rebalancing.right_guard;
                            let right_node_next = &mut rebalancing.right_right_guard;
                            let left_pos = rebalancing.parent_pos;
                            let right_pos = left_pos + 1;
                            if left_node.is_empty() || right_node.is_empty() {
                                return;
                            }
                            rebalancing.parent.innode_mut().merge_children(
                                left_pos,
                                right_pos,
                                left_node,
                                right_node,
                                right_node_next,
                                tree
                            );
                        },
                    ))
                // None
                } else if rebalancing.parent.len() == 1
                    && rebalancing.left_guard.is_empty()
                    && rebalancing.right_guard.is_empty()
                {
                    // this node and its children is empty, should be removed
                    debug!("Cleaning up empty level nodes, level {}", level);
                    Some(with_innode_removing(
                        key,
                        rebalancing,
                        &sub_node,
                        parent,
                        removed,
                        level,
                        |rebalancing| {
                            rebalancing.parent.remove(0);
                            rebalancing.parent.innode_mut().ptrs.as_slice()[0] = Default::default();
                        },
                    ))
                } else {
                    None
                }
            } else {
                None
            }
            .unwrap_or_else(|| RemoveResult {
                rebalancing: None,
                removed,
            })
        }
        MutSearchResult::External => {
            let node_guard: NodeWriteGuard<KS, PS> = write_node(node_ref);
            let mut target_guard = write_targeted(node_guard, key);
            let target_guard_ref = target_guard.node_ref().clone();
            let mut remove_result = RemoveResult {
                rebalancing: None,
                removed: false,
            };
            {
                let is_left_half_full = target_guard.is_half_full();
                let node = target_guard.extnode_mut(tree);
                let pos = node.search(key);
                let mut right_guard = write_node(&node.next);
                let _right_node_cannot_rebalance =
                    right_guard.is_none() || !right_guard.is_empty_node();
                if !right_guard.is_none() && !is_left_half_full || !right_guard.is_half_full() {
                    let right_right_guard =
                        write_node(&right_guard.right_ref_mut_no_empty().unwrap());
                    let parent_guard = write_node(parent);
                    let parent_target_guard = write_targeted(parent_guard, key);
                    let parent_pos = parent_target_guard.search(key);
                    // Check if the right node is innode and its parent is the same as the left one
                    // because we have to lock from left to right, there is no way to lock backwards
                    // if the empty non half-full node is at the right most of its parent
                    // So left over imbalanced such nodes will be expected and they can be eliminate
                    // by their left node remove operations.
                    if parent_pos < parent_target_guard.len() - 1 {
                        debug!("removing innode have a rebalance requirement");
                        remove_result.rebalancing = Some(RebalancingNodes {
                            left_ref: target_guard_ref,
                            left_guard: Default::default(),
                            parent: parent_target_guard,
                            parent_pos,
                            right_right_guard,
                            right_guard,
                        });
                    } else {
                        debug!("removing innode have a rebalance requirement but left-right node does not with the same parent");
                    }
                }
                if pos >= node.len {
                    debug!(
                        "Removing pos overflows external node, pos {}, len {}, expecting key {:?}, current keys {:?}, right keys {:?}",
                        pos, node.len, key, node.keys, read_unchecked::<KS, PS>(&node.next).keys()
                    );
                    remove_result.removed = false;
                }
                if &node.keys.as_slice()[pos] == key {
                    debug!("Removing key {:?} at {}, keys {:?}", key, pos, &node.keys);
                    node.remove_at(pos);
                    remove_result.removed = true;
                } else {
                    debug!(
                        "Search check failed for remove at pos {}, expecting {:?}, actual {:?}, have {:?}",
                        pos, key, &node.keys.as_slice_immute()[pos], &node.keys
                    );
                    remove_result.removed = false;
                }
            }
            if let Some(ref mut rebalance) = &mut remove_result.rebalancing {
                rebalance.left_guard = target_guard;
            }
            remove_result
        }
    }
}

// scatter the node and its references to ensure garbage collection
pub fn scatter_nodes<KS, PS>(node_ref: &NodeCellRef) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let mut footprint = 0;
    if node_ref.is_default() {
        return footprint;
    }
    let mut node = write_node::<KS, PS>(node_ref);
    loop {
        match &*node {
            &NodeData::None => return footprint,
            &NodeData::External(ref node) => {
                external::make_deleted(&node.id);
            }
            _ => {}
        }
        let right = node.right_ref().unwrap().clone();
        footprint += node.len();
        *node = NodeData::None;
        node = write_node::<KS, PS>(&right);
    }
}
<commit_msg>cleanup<commit_after>use crate::index::btree::node::write_node;
use crate::index::btree::node::NodeData;
use crate::index::btree::NodeCellRef;
use crate::index::btree::{external};
use crate::index::trees::EntryKey;
use crate::index::trees::Slice;
use std::fmt::Debug;

// scatter the node and its references to ensure garbage collection
pub fn scatter_nodes<KS, PS>(node_ref: &NodeCellRef) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let mut footprint = 0;
    if node_ref.is_default() {
        return footprint;
    }
    let mut node = write_node::<KS, PS>(node_ref);
    loop {
        match &*node {
            &NodeData::None => return footprint,
            &NodeData::External(ref node) => {
                external::make_deleted(&node.id);
            }
            _ => {}
        }
        let right = node.right_ref().unwrap().clone();
        footprint += node.len();
        *node = NodeData::None;
        node = write_node::<KS, PS>(&right);
    }
}
<|endoftext|>


<commit_before>extern crate time;
extern crate serde;
extern crate serde_json;
extern crate mongodb;
extern crate bson;

use super::worker_pool;
use super::cove;

use std::thread;
use std::sync::{Arc, Mutex};
use self::serde_json::Value;
use self::mongodb::db::{Database, ThreadedDatabase};
use self::mongodb::coll::options::FindOneAndUpdateOptions;
use self::bson::Document;
use self::bson::Bson;

///
/// Holds a program object
///
pub struct Program {
    pub data: String,
    pub program_id: u64
}

///
/// Does the actual ingestion
///
pub fn ingest(first_time: bool, db: &Database, num_workers: usize) {
    let total_programs = get_total_programs();
    let mut updated_date = Arc::new(String::from(""));

    if !first_time {
        let current_time = time::now_utc();
        updated_date = Arc::new(format!("{}-{}-{}", 1900 + current_time.tm_year, current_time.tm_mon, current_time.tm_mday));
    }

    let mut worker_pool = worker_pool::WorkerPool::new(num_workers);

    for i in (0..total_programs).step_by(200) {
        let shared_updated_date = updated_date.clone();
        let shared_db = db.clone();

        worker_pool.add_worker(thread::spawn(move || {
            let programs = get_programs(i);
            let mut worker_pool = worker_pool::WorkerPool::new(num_workers);
            let coll = shared_db.collection("programs");

            for program in programs {
                let video_ids = Arc::new(Mutex::new(vec![]));
                let program_id = program.program_id;

                let filter = doc! {
                    "program_id" => program_id
                };

                // Can't use doc! macro because it escapes the JSON data.
                let mut update = Document::new();
                update.insert("program_id", program_id);
                update.insert("data", Bson::JavaScriptCode(program.data.clone()));

                let mut options = FindOneAndUpdateOptions::new();
                options.upsert = true;
                coll.find_one_and_replace(filter, update, Some(options)).expect("Can't insert program into database!");

                let shared_updated_date = shared_updated_date.clone();
                let shared_db = shared_db.clone();

                // We need this final shared DB because shared_db is moved into the thread so it can't be used to delete
                // at the end.
                let final_shared_db = shared_db.clone();

                let shared_video_ids = video_ids.clone();

                worker_pool.add_worker(thread::spawn(move || {
                    let total_videos = get_video_count_for_program(&shared_updated_date, &program);
                    let mut worker_pool = worker_pool::WorkerPool::new(num_workers);
                    let program = Arc::new(program);
                    let shared_video_ids = shared_video_ids.clone();

                    for j in (0..total_videos).step_by(200) {
                        let shared_program = program.clone();
                        let shared_updated_date = shared_updated_date.clone();
                        let shared_db = shared_db.clone();
                        let shared_video_ids = shared_video_ids.clone();

                        worker_pool.add_worker(thread::spawn(move || {
                            get_videos(&shared_updated_date, j, &shared_program, &shared_db, &shared_video_ids);
                        }));
                    }

                    worker_pool.wait_for_children();
                }));

                let video_ids_str = format!("[{}]", video_ids.lock().unwrap().join(","));
                let filter = doc! {
                    "tp_media_object_id" => video_ids_str,
                    "cond" => "$nin"
                };

                let video_coll = final_shared_db.collection("videos");
                video_coll.delete_many(filter, None).unwrap();
            }

            worker_pool.wait_for_children();
        }));
    }

    worker_pool.wait_for_children();
}

///
/// Gets the total programs to break them up
///
fn get_total_programs() -> u64 {
    cove::video_api("programs", vec![]).as_object().unwrap().get("count").unwrap().as_u64().unwrap()
}

///
/// Gets all the shows from COVE
///
fn get_programs(start_index: u64) -> Vec<Program> {
    let cove_data = cove::video_api("programs", vec![["limit_start", start_index.to_string().as_str()]]);
    let programs: &Vec<Value> = cove_data.as_object().unwrap().get("results").unwrap().as_array().unwrap();
    let mut programs_data = vec![];

    for program in programs {
        let program_uri = program.as_object().unwrap().get("resource_uri").unwrap().as_str().unwrap();
        let program_id: u64 = program_uri.split("/").nth(4).unwrap().parse().unwrap();

        programs_data.push(Program {data: program.to_string(), program_id: program_id});
    }

    programs_data
}

///
/// Gets the total videos for a program so they can be chunked
///
fn get_video_count_for_program<'a>(updated_date: &Arc<String>, program: &Program) -> u64 {
    let program_id = program.program_id.to_string();
    let mut params = vec![
        ["filter_program", program_id.as_str()],
        ["exclude_type", "Other"]
    ];

    if updated_date.as_str() != "" {
        params.push(["filter_record_last_updated_datetime__gt", updated_date.as_str()]);
    }

    cove::video_api("videos", params).as_object().unwrap().get("count").unwrap().as_u64().unwrap()
}

///
/// Gets all videos from COVE for a program, 200 at a time
///
fn get_videos(updated_date: &Arc<String>, start_index: u64, program: &Arc<Program>, db: &Database, video_ids: &Arc<Mutex<Vec<String>>>) {
    let program_id = program.program_id.to_string();
    let str_start_index = start_index.to_string();
    let mut params = vec![
        ["filter_program", program_id.as_str()],
        ["limit_start", str_start_index.as_str()],
        ["exclude_type", "Other"]
    ];

    if updated_date.as_str() != "" {
        params.push(["filter_record_last_updated_datetime__gt", updated_date.as_str()]);
    }

    let cove_data = cove::video_api("videos", params);
    let videos:&Vec<Value> = cove_data.as_object().unwrap().get("results").unwrap().as_array().unwrap();
    let coll = db.collection("videos");
    let mut page_video_ids = vec![];

    for video in videos {

        //
        // Not everything has a tp_media_object_id.  Ignore the ones that don't.
        //
        match video.as_object().unwrap().get("tp_media_object_id").unwrap().as_u64() {
            Some(tp_media_object_id) => {
                page_video_ids.push(tp_media_object_id.to_string());

                let filter = doc! {
                    "tp_media_object_id" => tp_media_object_id
                };

                // Can't use doc! macro because it escapes the data.
                let mut update = Document::new();
                update.insert("tp_media_object_id", tp_media_object_id);
                update.insert("program_id", program.program_id);
                update.insert("data", Bson::JavaScriptCode(video.to_string()));

                let mut options = FindOneAndUpdateOptions::new();
                options.upsert = true;

                coll.find_one_and_replace(filter, update, Some(options)).expect("Can't insert video into database!");
            },
            None => {}
        };
    }

    let shared_video_ids = video_ids.clone();
    let mut locked_video_ids = shared_video_ids.lock().unwrap();
    locked_video_ids.extend(page_video_ids);
}<commit_msg>Fixed delete of old videos<commit_after>extern crate time;
extern crate serde;
extern crate serde_json;
extern crate mongodb;
extern crate bson;

use super::worker_pool;
use super::cove;

use std::thread;
use std::sync::{Arc, Mutex};
use self::serde_json::Value;
use self::mongodb::db::{Database, ThreadedDatabase};
use self::mongodb::coll::options::FindOneAndUpdateOptions;
use self::bson::{Document, Bson};

///
/// Holds a program object
///
pub struct Program {
    pub data: String,
    pub program_id: u64
}

///
/// Does the actual ingestion
///
pub fn ingest(first_time: bool, db: &Database, num_workers: usize) {
    let total_programs = get_total_programs();
    let mut updated_date = Arc::new(String::from(""));

    if !first_time {
        let current_time = time::now_utc();
        updated_date = Arc::new(format!("{}-{}-{}", 1900 + current_time.tm_year, current_time.tm_mon, current_time.tm_mday));
    }

    let mut worker_pool = worker_pool::WorkerPool::new(num_workers);

    for i in (0..total_programs).step_by(200) {
        let shared_updated_date = updated_date.clone();
        let shared_db = db.clone();

        worker_pool.add_worker(thread::spawn(move || {
            let programs = get_programs(i);
            let mut worker_pool = worker_pool::WorkerPool::new(num_workers);
            let coll = shared_db.collection("programs");

            for program in programs {
                let video_ids = Arc::new(Mutex::new(vec![]));
                let program_id = program.program_id;

                let filter = doc! {
                    "program_id" => program_id
                };

                // Can't use doc! macro because it escapes the JSON data.
                let mut update = Document::new();
                update.insert("program_id", program_id);
                update.insert("data", Bson::JavaScriptCode(program.data.clone()));

                let mut options = FindOneAndUpdateOptions::new();
                options.upsert = true;
                coll.find_one_and_replace(filter, update, Some(options)).expect("Can't insert program into database!");

                let shared_updated_date = shared_updated_date.clone();
                let shared_db = shared_db.clone();
                let shared_video_ids = video_ids.clone();

                worker_pool.add_worker(thread::spawn(move || {
                    let total_videos = get_video_count_for_program(&shared_updated_date, &program);
                    let mut worker_pool = worker_pool::WorkerPool::new(num_workers);
                    let program = Arc::new(program);
                    let shared_video_ids = shared_video_ids.clone();
                    let final_shared_video_ids = shared_video_ids.clone();

                    for j in (0..total_videos).step_by(200) {
                        let shared_program = program.clone();
                        let shared_updated_date = shared_updated_date.clone();
                        let shared_db = shared_db.clone();
                        let shared_video_ids = shared_video_ids.clone();

                        worker_pool.add_worker(thread::spawn(move || {
                            get_videos(&shared_updated_date, j, &shared_program, &shared_db, &shared_video_ids);
                        }));
                    }

                    worker_pool.wait_for_children();

                    // Delete videos that no longer exist.
                    let video_ids_lock = final_shared_video_ids.lock().unwrap();
                    let video_ids_arr = Bson::Array(video_ids_lock.clone().iter().map(|x| bson::to_bson(x).unwrap()).collect());

                    let filter = doc! {
                        "$and" => [
                        {
                            "program_id" => {
                                "$eq" => program_id
                            }
                        },
                        {
                            "tp_media_object_id" => {
                                "$nin" => video_ids_arr
                            }
                        }
                    ]};

                    let video_coll = shared_db.collection("videos");
                    video_coll.delete_many(filter, None).ok().expect("Error deleting records from database!");
                }));
            }

            worker_pool.wait_for_children();
        }));
    }

    worker_pool.wait_for_children();
}

///
/// Gets the total programs to break them up
///
fn get_total_programs() -> u64 {
    cove::video_api("programs", vec![]).as_object().unwrap().get("count").unwrap().as_u64().unwrap()
}

///
/// Gets all the shows from COVE
///
fn get_programs(start_index: u64) -> Vec<Program> {
    let cove_data = cove::video_api("programs", vec![["limit_start", start_index.to_string().as_str()]]);
    let programs: &Vec<Value> = cove_data.as_object().unwrap().get("results").unwrap().as_array().unwrap();
    let mut programs_data = vec![];

    for program in programs {
        let program_uri = program.as_object().unwrap().get("resource_uri").unwrap().as_str().unwrap();
        let program_id: u64 = program_uri.split("/").nth(4).unwrap().parse().unwrap();

        programs_data.push(Program {data: program.to_string(), program_id: program_id});
    }

    programs_data
}

///
/// Gets the total videos for a program so they can be chunked
///
fn get_video_count_for_program<'a>(updated_date: &Arc<String>, program: &Program) -> u64 {
    let program_id = program.program_id.to_string();
    let mut params = vec![
        ["filter_program", program_id.as_str()],
        ["exclude_type", "Other"]
    ];

    if updated_date.as_str() != "" {
        params.push(["filter_record_last_updated_datetime__gt", updated_date.as_str()]);
    }

    cove::video_api("videos", params).as_object().unwrap().get("count").unwrap().as_u64().unwrap()
}

///
/// Gets all videos from COVE for a program, 200 at a time
///
fn get_videos(updated_date: &Arc<String>, start_index: u64, program: &Arc<Program>, db: &Database, video_ids: &Arc<Mutex<Vec<u64>>>) {
    let program_id = program.program_id.to_string();
    let str_start_index = start_index.to_string();
    let mut params = vec![
        ["filter_program", program_id.as_str()],
        ["limit_start", str_start_index.as_str()],
        ["exclude_type", "Other"]
    ];

    if updated_date.as_str() != "" {
        params.push(["filter_record_last_updated_datetime__gt", updated_date.as_str()]);
    }

    let cove_data = cove::video_api("videos", params);
    let videos:&Vec<Value> = cove_data.as_object().unwrap().get("results").unwrap().as_array().unwrap();
    let coll = db.collection("videos");
    let mut page_video_ids = vec![];

    for video in videos {

        //
        // Not everything has a tp_media_object_id.  Ignore the ones that don't.
        //
        match video.as_object().unwrap().get("tp_media_object_id").unwrap().as_u64() {
            Some(tp_media_object_id) => {
                page_video_ids.push(tp_media_object_id);

                let filter = doc! {
                    "tp_media_object_id" => tp_media_object_id
                };

                // Can't use doc! macro because it escapes the data.
                let mut update = Document::new();
                update.insert("tp_media_object_id", tp_media_object_id);
                update.insert("program_id", program.program_id);
                update.insert("data", Bson::JavaScriptCode(video.to_string()));

                let mut options = FindOneAndUpdateOptions::new();
                options.upsert = true;

                coll.find_one_and_replace(filter, update, Some(options)).expect("Can't insert video into database!");
            },
            None => {}
        };
    }

    let shared_video_ids = video_ids.clone();
    let mut locked_video_ids = shared_video_ids.lock().unwrap();
    locked_video_ids.extend(page_video_ids);
}<|endoftext|>


<commit_before>use rustc::hir::*;
use rustc::lint::*;
use utils::{get_trait_def_id, implements_trait, higher, match_qpath, paths, span_lint};

/// **What it does:** Checks for iteration that is guaranteed to be infinite.
///
/// **Why is this bad?** While there may be places where this is acceptable
/// (e.g. in event streams), in most cases this is simply an error.
///
/// **Known problems:** None.
///
/// **Example:**
/// ```rust
/// repeat(1_u8).iter().collect::<Vec<_>>()
/// ```
declare_lint! {
    pub INFINITE_ITER,
    Warn,
    "infinite iteration"
}

/// **What it does:** Checks for iteration that may be infinite.
///
/// **Why is this bad?** While there may be places where this is acceptable
/// (e.g. in event streams), in most cases this is simply an error.
///
/// **Known problems:** The code may have a condition to stop iteration, but
/// this lint is not clever enough to analyze it.
///
/// **Example:**
/// ```rust
/// [0..].iter().zip(infinite_iter.take_while(|x| x > 5))
/// ```
declare_lint! {
    pub MAYBE_INFINITE_ITER,
    Allow,
    "possible infinite iteration"
}

#[derive(Copy, Clone)]
pub struct Pass;

impl LintPass for Pass {
    fn get_lints(&self) -> LintArray {
        lint_array!(INFINITE_ITER, MAYBE_INFINITE_ITER)
    }
}

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for Pass {
    fn check_expr(&mut self, cx: &LateContext<'a, 'tcx>, expr: &'tcx Expr) {
        let (lint, msg) = match complete_infinite_iter(cx, expr) {
            True => (INFINITE_ITER, "infinite iteration detected"),
            Unknown => (MAYBE_INFINITE_ITER,
                        "possible infinite iteration detected"),
            False => { return; }
        };
        span_lint(cx, lint, expr.span, msg)
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum TriState {
    True,
    Unknown,
    False
}

use self::TriState::{True, Unknown, False};

impl TriState {
    fn and(self, b: Self) -> Self {
        match (self, b) {
            (False, _) | (_, False) => False,
            (Unknown, _) | (_, Unknown) => Unknown,
            _ => True
        }
    }

    fn or(self, b: Self) -> Self {
        match (self, b) {
            (True, _) | (_, True) => True,
            (Unknown, _) | (_, Unknown) => Unknown,
            _ => False
        }
    }
}

impl From<bool> for TriState {
    fn from(b: bool) -> Self {
        if b { True } else { False }
    }
}

#[derive(Copy, Clone)]
enum Heuristic {
    Always,
    First,
    Any,
    All
}

use self::Heuristic::{Always, First, Any, All};

// here we use the `TriState` as (Finite, Possible Infinite, Infinite)
static HEURISTICS : &[(&str, usize, Heuristic, TriState)] = &[
    ("zip", 2, All, True),
    ("chain", 2, Any, True),
    ("cycle", 1, Always, True),
    ("map", 2, First, True),
    ("by_ref", 1, First, True),
    ("cloned", 1, First, True),
    ("rev", 1, First, True),
    ("inspect", 1, First, True),
    ("enumerate", 1, First, True),
    ("peekable", 2, First, True),
    ("fuse", 1, First, True),
    ("skip", 2, First, True),
    ("skip_while", 1, First, True),
    ("filter", 2, First, True),
    ("filter_map", 2, First, True),
    ("flat_map", 2, First, True),
    ("unzip", 1, First, True),
    ("take_while", 2, First, Unknown),
    ("scan", 3, First, Unknown)
];

fn is_infinite(cx: &LateContext, expr: &Expr) -> TriState {
    match expr.node {
        ExprMethodCall(ref method, _, ref args) => {
            for &(name, len, heuristic, cap) in HEURISTICS.iter() {
                if method.name == name && args.len() == len {
                    return (match heuristic {
                        Always => True,
                        First => is_infinite(cx, &args[0]),
                        Any => is_infinite(cx, &args[0]).or(is_infinite(cx, &args[1])),
                        All => is_infinite(cx, &args[0]).and(is_infinite(cx, &args[1])),
                    }).and(cap);
                }
            }
            if method.name == "flat_map" && args.len() == 2 {
                if let ExprClosure(_, _, body_id, _) = args[1].node {
                    let body = cx.tcx.hir.body(body_id);
                    return is_infinite(cx, &body.value);
                }
            }
            False
        },
        ExprBlock(ref block) =>
            block.expr.as_ref().map_or(False, |e| is_infinite(cx, e)),
        ExprBox(ref e) | ExprAddrOf(_, ref e) => is_infinite(cx, e),
        ExprCall(ref path, _) => {
            if let ExprPath(ref qpath) = path.node {
                match_qpath(qpath, &paths::REPEAT).into()
            } else { False }
        },
        ExprStruct(..) => {
            higher::range(expr).map_or(false, |r| r.end.is_none()).into()
        },
        _ => False
    }
}

static POSSIBLY_COMPLETING_METHODS : &[(&str, usize)] = &[
    ("find", 2),
    ("rfind", 2),
    ("position", 2),
    ("rposition", 2),
    ("any", 2),
    ("all", 2)
];

static COMPLETING_METHODS : &[(&str, usize)] = &[
    ("count", 1),
    ("collect", 1),
    ("fold", 3),
    ("for_each", 2),
    ("partition", 2),
    ("max", 1),
    ("max_by", 2),
    ("max_by_key", 2),
    ("min", 1),
    ("min_by", 2),
    ("min_by_key", 2),
    ("sum", 1),
    ("product", 1)
];

fn complete_infinite_iter(cx: &LateContext, expr: &Expr) -> TriState {
    match expr.node {
        ExprMethodCall(ref method, _, ref args) => {
            for &(name, len) in COMPLETING_METHODS.iter() {
                if method.name == name && args.len() == len {
                    return is_infinite(cx, &args[0]);
                }
            }
            for &(name, len) in POSSIBLY_COMPLETING_METHODS.iter() {
                if method.name == name && args.len() == len {
                    return Unknown.and(is_infinite(cx, &args[0]));
                }
            }
            if method.name == "last" && args.len() == 1 &&
                    get_trait_def_id(cx, &paths::DOUBLE_ENDED_ITERATOR).map_or(false,
                        |id| !implements_trait(cx,
                                               cx.tables.expr_ty(&args[0]),
                                               id,
                                               &[])) {
                return is_infinite(cx, &args[0]);
            }
        },
        ExprBinary(op, ref l, ref r) => {
            if op.node.is_comparison() {
                return is_infinite(cx, l).and(is_infinite(cx, r)).and(Unknown)
            }
        }, //TODO: ExprLoop + Match
        _ => ()
    }
    False
}
<commit_msg>rename TriState -> Finiteness, docs<commit_after>use rustc::hir::*;
use rustc::lint::*;
use utils::{get_trait_def_id, implements_trait, higher, match_qpath, paths, span_lint};

/// **What it does:** Checks for iteration that is guaranteed to be infinite.
///
/// **Why is this bad?** While there may be places where this is acceptable
/// (e.g. in event streams), in most cases this is simply an error.
///
/// **Known problems:** None.
///
/// **Example:**
/// ```rust
/// repeat(1_u8).iter().collect::<Vec<_>>()
/// ```
declare_lint! {
    pub INFINITE_ITER,
    Warn,
    "infinite iteration"
}

/// **What it does:** Checks for iteration that may be infinite.
///
/// **Why is this bad?** While there may be places where this is acceptable
/// (e.g. in event streams), in most cases this is simply an error.
///
/// **Known problems:** The code may have a condition to stop iteration, but
/// this lint is not clever enough to analyze it.
///
/// **Example:**
/// ```rust
/// [0..].iter().zip(infinite_iter.take_while(|x| x > 5))
/// ```
declare_lint! {
    pub MAYBE_INFINITE_ITER,
    Allow,
    "possible infinite iteration"
}

#[derive(Copy, Clone)]
pub struct Pass;

impl LintPass for Pass {
    fn get_lints(&self) -> LintArray {
        lint_array!(INFINITE_ITER, MAYBE_INFINITE_ITER)
    }
}

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for Pass {
    fn check_expr(&mut self, cx: &LateContext<'a, 'tcx>, expr: &'tcx Expr) {
        let (lint, msg) = match complete_infinite_iter(cx, expr) {
            Infinite => (INFINITE_ITER, "infinite iteration detected"),
            MaybeInfinite => (MAYBE_INFINITE_ITER,
                        "possible infinite iteration detected"),
            Finite => { return; }
        };
        span_lint(cx, lint, expr.span, msg)
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum Finiteness {
    Infinite,
    MaybeInfinite,
    Finite
}

use self::Finiteness::{Infinite, MaybeInfinite, Finite};

impl Finiteness {
    fn and(self, b: Self) -> Self {
        match (self, b) {
            (Finite, _) | (_, Finite) => Finite,
            (MaybeInfinite, _) | (_, MaybeInfinite) => MaybeInfinite,
            _ => Infinite
        }
    }

    fn or(self, b: Self) -> Self {
        match (self, b) {
            (Infinite, _) | (_, Infinite) => Infinite,
            (MaybeInfinite, _) | (_, MaybeInfinite) => MaybeInfinite,
            _ => Finite
        }
    }
}

impl From<bool> for Finiteness {
    fn from(b: bool) -> Self {
        if b { Infinite } else { Finite }
    }
}

/// This tells us what to look for to know if the iterator returned by
/// this method is infinite
#[derive(Copy, Clone)]
enum Heuristic {
    /// infinite no matter what
    Always,
    /// infinite if the first argument is
    First,
    /// infinite if any of the supplied arguments is
    Any,
    /// infinite if all of the supplied arguments are
    All
}

use self::Heuristic::{Always, First, Any, All};

/// a slice of (method name, number of args, heuristic, bounds) tuples
/// that will be used to determine whether the method in question
/// returns an infinite or possibly infinite iterator. The finiteness
/// is an upper bound, e.g. some methods can return a possibly
/// infinite iterator at worst, e.g. `take_while`.
static HEURISTICS : &[(&str, usize, Heuristic, Finiteness)] = &[
    ("zip", 2, All, Infinite),
    ("chain", 2, Any, Infinite),
    ("cycle", 1, Always, Infinite),
    ("map", 2, First, Infinite),
    ("by_ref", 1, First, Infinite),
    ("cloned", 1, First, Infinite),
    ("rev", 1, First, Infinite),
    ("inspect", 1, First, Infinite),
    ("enumerate", 1, First, Infinite),
    ("peekable", 2, First, Infinite),
    ("fuse", 1, First, Infinite),
    ("skip", 2, First, Infinite),
    ("skip_while", 1, First, Infinite),
    ("filter", 2, First, Infinite),
    ("filter_map", 2, First, Infinite),
    ("flat_map", 2, First, Infinite),
    ("unzip", 1, First, Infinite),
    ("take_while", 2, First, MaybeInfinite),
    ("scan", 3, First, MaybeInfinite)
];

fn is_infinite(cx: &LateContext, expr: &Expr) -> Finiteness {
    match expr.node {
        ExprMethodCall(ref method, _, ref args) => {
            for &(name, len, heuristic, cap) in HEURISTICS.iter() {
                if method.name == name && args.len() == len {
                    return (match heuristic {
                        Always => Infinite,
                        First => is_infinite(cx, &args[0]),
                        Any => is_infinite(cx, &args[0]).or(is_infinite(cx, &args[1])),
                        All => is_infinite(cx, &args[0]).and(is_infinite(cx, &args[1])),
                    }).and(cap);
                }
            }
            if method.name == "flat_map" && args.len() == 2 {
                if let ExprClosure(_, _, body_id, _) = args[1].node {
                    let body = cx.tcx.hir.body(body_id);
                    return is_infinite(cx, &body.value);
                }
            }
            Finite
        },
        ExprBlock(ref block) =>
            block.expr.as_ref().map_or(Finite, |e| is_infinite(cx, e)),
        ExprBox(ref e) | ExprAddrOf(_, ref e) => is_infinite(cx, e),
        ExprCall(ref path, _) => {
            if let ExprPath(ref qpath) = path.node {
                match_qpath(qpath, &paths::REPEAT).into()
            } else { Finite }
        },
        ExprStruct(..) => {
            higher::range(expr).map_or(false, |r| r.end.is_none()).into()
        },
        _ => Finite
    }
}

/// the names and argument lengths of methods that *may* exhaust their
/// iterators
static POSSIBLY_COMPLETING_METHODS : &[(&str, usize)] = &[
    ("find", 2),
    ("rfind", 2),
    ("position", 2),
    ("rposition", 2),
    ("any", 2),
    ("all", 2)
];

/// the names and argument lengths of methods that *always* exhaust
/// their iterators
static COMPLETING_METHODS : &[(&str, usize)] = &[
    ("count", 1),
    ("collect", 1),
    ("fold", 3),
    ("for_each", 2),
    ("partition", 2),
    ("max", 1),
    ("max_by", 2),
    ("max_by_key", 2),
    ("min", 1),
    ("min_by", 2),
    ("min_by_key", 2),
    ("sum", 1),
    ("product", 1)
];

fn complete_infinite_iter(cx: &LateContext, expr: &Expr) -> Finiteness {
    match expr.node {
        ExprMethodCall(ref method, _, ref args) => {
            for &(name, len) in COMPLETING_METHODS.iter() {
                if method.name == name && args.len() == len {
                    return is_infinite(cx, &args[0]);
                }
            }
            for &(name, len) in POSSIBLY_COMPLETING_METHODS.iter() {
                if method.name == name && args.len() == len {
                    return MaybeInfinite.and(is_infinite(cx, &args[0]));
                }
            }
            if method.name == "last" && args.len() == 1 &&
                    get_trait_def_id(cx, &paths::DOUBLE_ENDED_ITERATOR).map_or(false,
                        |id| !implements_trait(cx,
                                               cx.tables.expr_ty(&args[0]),
                                               id,
                                               &[])) {
                return is_infinite(cx, &args[0]);
            }
        },
        ExprBinary(op, ref l, ref r) => {
            if op.node.is_comparison() {
                return is_infinite(cx, l).and(is_infinite(cx, r)).and(MaybeInfinite)
            }
        }, //TODO: ExprLoop + Match
        _ => ()
    }
    Finite
}
<|endoftext|>


<commit_before>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use mundane::hash::{Digest, Hasher, Sha256};
use std::time::Duration;
use thiserror::Error;
use zerocopy::AsBytes;

use crate::app::TransportWrapper;
use crate::bootstrap::{Bootstrap, BootstrapOptions, UpdateProtocol};
use crate::io::spi::Transfer;
use crate::transport::{Capability, Result};

#[derive(AsBytes, Debug, Default)]
#[repr(C)]
struct FrameHeader {
    hash: [u8; Frame::HASH_LEN],
    frame_num: u32,
    flash_offset: u32,
}

#[derive(AsBytes, Debug)]
#[repr(C)]
struct Frame {
    header: FrameHeader,
    data: [u8; Frame::DATA_LEN],
}

impl Default for Frame {
    fn default() -> Self {
        Frame {
            header: Default::default(),
            data: [0xff; Frame::DATA_LEN],
        }
    }
}

impl Frame {
    const EOF: u32 = 0x8000_0000;
    const FLASH_SECTOR_SIZE: usize = 2048;
    const FLASH_SECTOR_MASK: usize = Self::FLASH_SECTOR_SIZE - 1;
    const FLASH_BUFFER_SIZE: usize = 128;
    const FLASH_BUFFER_MASK: usize = Self::FLASH_BUFFER_SIZE - 1;
    const DATA_LEN: usize = 2048 - std::mem::size_of::<FrameHeader>();
    const HASH_LEN: usize = 32;
    const FLASH_BASE_ADDRESS: usize = 65536 * 8;

    /// Computes the hash in the header.
    fn header_hash(&self) -> [u8; Frame::HASH_LEN] {
        let frame = self.as_bytes();
        let sha = Sha256::hash(&frame[Frame::HASH_LEN..]);
        sha.bytes()
    }

    /// Computes the hash over the entire frame.
    fn frame_hash(&self) -> [u8; Frame::HASH_LEN] {
        let sha = Sha256::hash(self.as_bytes());
        let mut digest = sha.bytes();
        // Touch up zeroes into ones, as that is what the old chips are doing.
        for b in &mut digest {
            if *b == 0 {
                *b = 1;
            }
        }
        digest
    }

    /// Creates a sequence of frames based on a `payload` binary.
    fn from_payload(payload: &[u8]) -> Vec<Frame> {
        let mut frames = Vec::new();

        let max_addr = (payload
            .chunks(4)
            .rposition(|c| c != &[0xff; 4])
            .unwrap_or(0)
            + 1)
            * 4;

        let mut frame_num = 0;
        let mut addr = 0;
        while addr < max_addr {
            // Try skipping over 0xffffffff words.
            let nonempty_addr = addr
                + payload[addr..]
                    .chunks(4)
                    .position(|c| c != &[0xff; 4])
                    .unwrap()
                    * 4;
            let skip_addr = nonempty_addr & !Self::FLASH_SECTOR_MASK;
            if skip_addr > addr && (addr == 0 || addr & Self::FLASH_BUFFER_MASK != 0) {
                // Can only skip from the start or if the last addr wasn't an exact multiple of
                // 128 (per H1D boot rom).
                addr = skip_addr;
            }

            let mut frame = Frame {
                header: FrameHeader {
                    frame_num,
                    flash_offset: (addr + Self::FLASH_BASE_ADDRESS) as u32,
                    ..Default::default()
                },
                ..Default::default()
            };
            let slice_size = Self::DATA_LEN.min(payload.len() - addr);
            frame.data[..slice_size].copy_from_slice(&payload[addr..addr + slice_size]);
            frames.push(frame);

            addr += Self::DATA_LEN;
            frame_num += 1;
        }
        frames.last_mut().map(|f| f.header.frame_num |= Self::EOF);
        frames
            .iter_mut()
            .for_each(|f| f.header.hash = f.header_hash());
        frames
    }
}

/// Implements the bootstrap protocol of previous Google Titan family chips.
pub struct Legacy {
    /// Controls the time that CS is deasserted between frames.  With Dauntless, 20ms has been
    /// observed to be "too long", in that after the final frame was transmitted, the bootloader
    /// could have already stopped listening to SPI and started booting the newly flashed code,
    /// before OpenTitan tool could read the ACK of the final transaction.  1ms is what the
    /// spiflash.cc tool from cr50-utils uses.
    pub inter_frame_delay: Duration,
}

impl Legacy {
    const INTER_FRAME_DELAY: Duration = Duration::from_millis(1);
    const MAX_CONSECUTIVE_ERRORS: u32 = 100;

    /// Creates a new `Primitive` protocol updater from `options`.
    pub fn new(options: &BootstrapOptions) -> Self {
        Self {
            inter_frame_delay: options.inter_frame_delay.unwrap_or(Self::INTER_FRAME_DELAY),
        }
    }
}

#[derive(Debug, Error, serde::Serialize, serde::Deserialize)]
pub enum LegacyBootstrapError {
    #[error("Boot rom not ready")]
    NotReady,
    #[error("Unknown boot rom error: {0}")]
    Unknown(u8),
    #[error("Boot rom error: NOREQUEST")]
    NoRequest,
    #[error("Boot rom error: NOMAGIC")]
    NoMagic,
    #[error("Boot rom error: TOOBIG")]
    TooBig,
    #[error("Boot rom error: TOOHIGH")]
    TooHigh,
    #[error("Boot rom error: NOALIGN")]
    NoAlign,
    #[error("Boot rom error: NOROUND")]
    NoRound,
    #[error("Boot rom error: BADKEY")]
    BadKey,
    #[error("Boot rom error: BADSTART")]
    BadStart,
    #[error("Boot rom error: NOWIPE")]
    NoWipe,
    #[error("Boot rom error: NOWIPE0")]
    NoWipe0,
    #[error("Boot rom error: NOWIPE1")]
    NoWipe1,
    #[error("Boot rom error: NOTEMPTY")]
    NotEmpty,
    #[error("Boot rom error: NOWRITE")]
    NoWrite,
    #[error("Boot rom error: BADADR")]
    BadAdr,
    #[error("Boot rom error: OVERFLOW")]
    Overflow,
    #[error("Repeated errors communicating with boot rom")]
    RepeatedErrors,
}

impl From<u8> for LegacyBootstrapError {
    fn from(value: u8) -> LegacyBootstrapError {
        match value {
            // All zeroes or all ones means that the bootloader is not yet ready to respond.
            0 | 255 => LegacyBootstrapError::NotReady,
            // Other values represent particular errors.
            1 => LegacyBootstrapError::NoRequest,
            2 => LegacyBootstrapError::NoMagic,
            3 => LegacyBootstrapError::TooBig,
            4 => LegacyBootstrapError::TooHigh,
            5 => LegacyBootstrapError::NoAlign,
            6 => LegacyBootstrapError::NoRound,
            7 => LegacyBootstrapError::BadKey,
            8 => LegacyBootstrapError::BadStart,
            10 => LegacyBootstrapError::NoWipe,
            11 => LegacyBootstrapError::NoWipe0,
            12 => LegacyBootstrapError::NoWipe1,
            13 => LegacyBootstrapError::NotEmpty,
            14 => LegacyBootstrapError::NoWrite,
            15 => LegacyBootstrapError::BadAdr,
            16 => LegacyBootstrapError::Overflow,
            n => LegacyBootstrapError::Unknown(n),
        }
    }
}

impl UpdateProtocol for Legacy {
    fn verify_capabilities(
        &self,
        _container: &Bootstrap,
        transport: &TransportWrapper,
    ) -> Result<()> {
        transport
            .capabilities()?
            .request(Capability::GPIO | Capability::SPI)
            .ok()?;
        Ok(())
    }

    fn uses_common_bootstrap_reset(&self) -> bool {
        true
    }

    /// Performs the update protocol using the `transport` with the firmware `payload`.
    fn update(
        &self,
        container: &Bootstrap,
        transport: &TransportWrapper,
        payload: &[u8],
    ) -> Result<()> {
        let spi = container.spi_params.create(transport)?;

        let frames = Frame::from_payload(payload);

        // All frames up to but not including this index have been ack'ed by the bootloader.
        // Once this reaches frames.len(), the operation was successful.
        let mut first_unacked_index = 0;

        // Counts the number of repeated failures in getting a frame ack.
        let mut consecutive_errors = 0;

        // Set if the past transaction ack'ed the block sent before that one, while transmitting
        // the next.  If so, we heuristically assume that the block just sent will be ack'ed in
        // the upcoming bidirectional transaction.
        let mut optimistic = false;

        // Set if the past transaction ack'ed the block sent before that one, while
        // re-transmitting the same block.  If so, we do not care about the ack in the upcoming
        // bidirectional transaction, as it turns out it was not necessary to perform the most
        // recent re-transmission.
        let mut becoming_optimistic = false;

        loop {
            if consecutive_errors > Self::MAX_CONSECUTIVE_ERRORS {
                return Err(LegacyBootstrapError::RepeatedErrors.into());
            }

            let second_unacked_index = (first_unacked_index + 1).min(frames.len() - 1);
            let transmit_index = if optimistic {
                second_unacked_index
            } else {
                first_unacked_index
            };
            let frame = &frames[transmit_index];
            eprint!("{}.", transmit_index);
            std::thread::sleep(self.inter_frame_delay);

            // Write the frame and read back the ack of a previously transmitted frame.
            let mut response = [0u8; std::mem::size_of::<Frame>()];
            spi.run_transaction(&mut [Transfer::Both(frame.as_bytes(), &mut response)])?;

            if becoming_optimistic {
                optimistic = true;
                becoming_optimistic = false;
                continue;
            }

            if response.iter().all(|&x| x == response[0]) {
                // A response consisteing of all identical bytes is a status code.
                match LegacyBootstrapError::from(response[0]) {
                    LegacyBootstrapError::NotReady => {
                        consecutive_errors += 1;
                        continue; // Retry sending same frame.
                    }
                    error => return Err(error.into()),
                }
            }

            if response[..Frame::HASH_LEN] == frames[first_unacked_index].frame_hash() {
                first_unacked_index = first_unacked_index + 1;
            } else if response[..Frame::HASH_LEN] == frames[second_unacked_index].frame_hash() {
                first_unacked_index = second_unacked_index + 1;
            } else {
                consecutive_errors += 1;
                optimistic = false;
                continue;
            }

            consecutive_errors = 0;
            if first_unacked_index == frames.len() {
                // All frames acked, we are done.
                break;
            }
            if !optimistic {
                // Heuristic, become optimistic after second or later frame is ack'ed.
                becoming_optimistic = first_unacked_index > 1;
            }
        }
        eprintln!("success");
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const SIMPLE_BIN: &[u8; 2048] = include_bytes!("simple.bin");

    #[test]
    fn test_small_binary() -> Result<()> {
        let frames = Frame::from_payload(SIMPLE_BIN);

        assert_eq!(frames[0].header.frame_num, 0);
        assert_eq!(frames[0].header.flash_offset, 0x80000);
        assert_eq!(
            hex::encode(frames[0].header.hash),
            "4e31bfd8b3be32358f2235c0f241f3970de575fc6aca0564aa6bf30adaf33910"
        );

        assert_eq!(frames[1].header.frame_num, 0x8000_0001);
        assert_eq!(frames[1].header.flash_offset, 0x807d8);
        assert_eq!(
            hex::encode(frames[1].header.hash),
            "ff584c07bbb0a039934a660bd49b7812af8ee847d1e675d9aba71c11fab3cfcb"
        );
        Ok(())
    }
}
<commit_msg>[opentitantool] Correctly detect bootstrap errors<commit_after>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use mundane::hash::{Digest, Hasher, Sha256};
use std::time::Duration;
use thiserror::Error;
use zerocopy::AsBytes;

use crate::app::TransportWrapper;
use crate::bootstrap::{Bootstrap, BootstrapOptions, UpdateProtocol};
use crate::io::spi::Transfer;
use crate::transport::{Capability, Result};

#[derive(AsBytes, Debug, Default)]
#[repr(C)]
struct FrameHeader {
    hash: [u8; Frame::HASH_LEN],
    frame_num: u32,
    flash_offset: u32,
}

#[derive(AsBytes, Debug)]
#[repr(C)]
struct Frame {
    header: FrameHeader,
    data: [u8; Frame::DATA_LEN],
}

impl Default for Frame {
    fn default() -> Self {
        Frame {
            header: Default::default(),
            data: [0xff; Frame::DATA_LEN],
        }
    }
}

impl Frame {
    const EOF: u32 = 0x8000_0000;
    const FLASH_SECTOR_SIZE: usize = 2048;
    const FLASH_SECTOR_MASK: usize = Self::FLASH_SECTOR_SIZE - 1;
    const FLASH_BUFFER_SIZE: usize = 128;
    const FLASH_BUFFER_MASK: usize = Self::FLASH_BUFFER_SIZE - 1;
    const DATA_LEN: usize = 2048 - std::mem::size_of::<FrameHeader>();
    const HASH_LEN: usize = 32;
    const FLASH_BASE_ADDRESS: usize = 65536 * 8;

    /// Computes the hash in the header.
    fn header_hash(&self) -> [u8; Frame::HASH_LEN] {
        let frame = self.as_bytes();
        let sha = Sha256::hash(&frame[Frame::HASH_LEN..]);
        sha.bytes()
    }

    /// Computes the hash over the entire frame.
    fn frame_hash(&self) -> [u8; Frame::HASH_LEN] {
        let sha = Sha256::hash(self.as_bytes());
        let mut digest = sha.bytes();
        // Touch up zeroes into ones, as that is what the old chips are doing.
        for b in &mut digest {
            if *b == 0 {
                *b = 1;
            }
        }
        digest
    }

    /// Creates a sequence of frames based on a `payload` binary.
    fn from_payload(payload: &[u8]) -> Vec<Frame> {
        let mut frames = Vec::new();

        let max_addr = (payload
            .chunks(4)
            .rposition(|c| c != &[0xff; 4])
            .unwrap_or(0)
            + 1)
            * 4;

        let mut frame_num = 0;
        let mut addr = 0;
        while addr < max_addr {
            // Try skipping over 0xffffffff words.
            let nonempty_addr = addr
                + payload[addr..]
                    .chunks(4)
                    .position(|c| c != &[0xff; 4])
                    .unwrap()
                    * 4;
            let skip_addr = nonempty_addr & !Self::FLASH_SECTOR_MASK;
            if skip_addr > addr && (addr == 0 || addr & Self::FLASH_BUFFER_MASK != 0) {
                // Can only skip from the start or if the last addr wasn't an exact multiple of
                // 128 (per H1D boot rom).
                addr = skip_addr;
            }

            let mut frame = Frame {
                header: FrameHeader {
                    frame_num,
                    flash_offset: (addr + Self::FLASH_BASE_ADDRESS) as u32,
                    ..Default::default()
                },
                ..Default::default()
            };
            let slice_size = Self::DATA_LEN.min(payload.len() - addr);
            frame.data[..slice_size].copy_from_slice(&payload[addr..addr + slice_size]);
            frames.push(frame);

            addr += Self::DATA_LEN;
            frame_num += 1;
        }
        frames.last_mut().map(|f| f.header.frame_num |= Self::EOF);
        frames
            .iter_mut()
            .for_each(|f| f.header.hash = f.header_hash());
        frames
    }
}

/// Implements the bootstrap protocol of previous Google Titan family chips.
pub struct Legacy {
    /// Controls the time that CS is deasserted between frames.  With Dauntless, 20ms has been
    /// observed to be "too long", in that after the final frame was transmitted, the bootloader
    /// could have already stopped listening to SPI and started booting the newly flashed code,
    /// before OpenTitan tool could read the ACK of the final transaction.  1ms is what the
    /// spiflash.cc tool from cr50-utils uses.
    pub inter_frame_delay: Duration,
}

impl Legacy {
    const INTER_FRAME_DELAY: Duration = Duration::from_millis(1);
    const MAX_CONSECUTIVE_ERRORS: u32 = 100;

    /// Creates a new `Primitive` protocol updater from `options`.
    pub fn new(options: &BootstrapOptions) -> Self {
        Self {
            inter_frame_delay: options.inter_frame_delay.unwrap_or(Self::INTER_FRAME_DELAY),
        }
    }
}

#[derive(Debug, Error, serde::Serialize, serde::Deserialize)]
pub enum LegacyBootstrapError {
    #[error("Boot rom not ready")]
    NotReady,
    #[error("Unknown boot rom error: {0}")]
    Unknown(u8),
    #[error("Boot rom error: NOREQUEST")]
    NoRequest,
    #[error("Boot rom error: NOMAGIC")]
    NoMagic,
    #[error("Boot rom error: TOOBIG")]
    TooBig,
    #[error("Boot rom error: TOOHIGH")]
    TooHigh,
    #[error("Boot rom error: NOALIGN")]
    NoAlign,
    #[error("Boot rom error: NOROUND")]
    NoRound,
    #[error("Boot rom error: BADKEY")]
    BadKey,
    #[error("Boot rom error: BADSTART")]
    BadStart,
    #[error("Boot rom error: NOWIPE")]
    NoWipe,
    #[error("Boot rom error: NOWIPE0")]
    NoWipe0,
    #[error("Boot rom error: NOWIPE1")]
    NoWipe1,
    #[error("Boot rom error: NOTEMPTY")]
    NotEmpty,
    #[error("Boot rom error: NOWRITE")]
    NoWrite,
    #[error("Boot rom error: BADADR")]
    BadAdr,
    #[error("Boot rom error: OVERFLOW")]
    Overflow,
    #[error("Repeated errors communicating with boot rom")]
    RepeatedErrors,
}

impl From<u8> for LegacyBootstrapError {
    fn from(value: u8) -> LegacyBootstrapError {
        match value {
            // All zeroes or all ones means that the bootloader is not yet ready to respond.
            0 | 255 => LegacyBootstrapError::NotReady,
            // Other values represent particular errors.
            1 => LegacyBootstrapError::NoRequest,
            2 => LegacyBootstrapError::NoMagic,
            3 => LegacyBootstrapError::TooBig,
            4 => LegacyBootstrapError::TooHigh,
            5 => LegacyBootstrapError::NoAlign,
            6 => LegacyBootstrapError::NoRound,
            7 => LegacyBootstrapError::BadKey,
            8 => LegacyBootstrapError::BadStart,
            10 => LegacyBootstrapError::NoWipe,
            11 => LegacyBootstrapError::NoWipe0,
            12 => LegacyBootstrapError::NoWipe1,
            13 => LegacyBootstrapError::NotEmpty,
            14 => LegacyBootstrapError::NoWrite,
            15 => LegacyBootstrapError::BadAdr,
            16 => LegacyBootstrapError::Overflow,
            n => LegacyBootstrapError::Unknown(n),
        }
    }
}

impl UpdateProtocol for Legacy {
    fn verify_capabilities(
        &self,
        _container: &Bootstrap,
        transport: &TransportWrapper,
    ) -> Result<()> {
        transport
            .capabilities()?
            .request(Capability::GPIO | Capability::SPI)
            .ok()?;
        Ok(())
    }

    fn uses_common_bootstrap_reset(&self) -> bool {
        true
    }

    /// Performs the update protocol using the `transport` with the firmware `payload`.
    fn update(
        &self,
        container: &Bootstrap,
        transport: &TransportWrapper,
        payload: &[u8],
    ) -> Result<()> {
        let spi = container.spi_params.create(transport)?;

        let frames = Frame::from_payload(payload);

        // All frames up to but not including this index have been ack'ed by the bootloader.
        // Once this reaches frames.len(), the operation was successful.
        let mut first_unacked_index = 0;

        // Counts the number of repeated failures in getting a frame ack.
        let mut consecutive_errors = 0;

        // Set if the past transaction ack'ed the block sent before that one, while transmitting
        // the next.  If so, we heuristically assume that the block just sent will be ack'ed in
        // the upcoming bidirectional transaction.
        let mut optimistic = false;

        // Set if the past transaction ack'ed the block sent before that one, while
        // re-transmitting the same block.  If so, we do not care about the ack in the upcoming
        // bidirectional transaction, as it turns out it was not necessary to perform the most
        // recent re-transmission.
        let mut becoming_optimistic = false;

        loop {
            if consecutive_errors > Self::MAX_CONSECUTIVE_ERRORS {
                return Err(LegacyBootstrapError::RepeatedErrors.into());
            }

            let second_unacked_index = (first_unacked_index + 1).min(frames.len() - 1);
            let transmit_index = if optimistic {
                second_unacked_index
            } else {
                first_unacked_index
            };
            let frame = &frames[transmit_index];
            eprint!("{}.", transmit_index);
            std::thread::sleep(self.inter_frame_delay);

            // Write the frame and read back the ack of a previously transmitted frame.
            let mut response = [0u8; std::mem::size_of::<Frame>()];
            spi.run_transaction(&mut [Transfer::Both(frame.as_bytes(), &mut response)])?;

            if becoming_optimistic {
                optimistic = true;
                becoming_optimistic = false;
                continue;
            }

            if response[..Frame::HASH_LEN]
                .iter()
                .all(|&x| x == response[0])
            {
                // A response consisteing of all identical bytes is a status code.
                match LegacyBootstrapError::from(response[0]) {
                    LegacyBootstrapError::NotReady => {
                        consecutive_errors += 1;
                        continue; // Retry sending same frame.
                    }
                    error => return Err(error.into()),
                }
            }

            if response[..Frame::HASH_LEN] == frames[first_unacked_index].frame_hash() {
                first_unacked_index = first_unacked_index + 1;
            } else if response[..Frame::HASH_LEN] == frames[second_unacked_index].frame_hash() {
                first_unacked_index = second_unacked_index + 1;
            } else {
                consecutive_errors += 1;
                optimistic = false;
                continue;
            }

            consecutive_errors = 0;
            if first_unacked_index == frames.len() {
                // All frames acked, we are done.
                break;
            }
            if !optimistic {
                // Heuristic, become optimistic after second or later frame is ack'ed.
                becoming_optimistic = first_unacked_index > 1;
            }
        }
        eprintln!("success");
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const SIMPLE_BIN: &[u8; 2048] = include_bytes!("simple.bin");

    #[test]
    fn test_small_binary() -> Result<()> {
        let frames = Frame::from_payload(SIMPLE_BIN);

        assert_eq!(frames[0].header.frame_num, 0);
        assert_eq!(frames[0].header.flash_offset, 0x80000);
        assert_eq!(
            hex::encode(frames[0].header.hash),
            "4e31bfd8b3be32358f2235c0f241f3970de575fc6aca0564aa6bf30adaf33910"
        );

        assert_eq!(frames[1].header.frame_num, 0x8000_0001);
        assert_eq!(frames[1].header.flash_offset, 0x807d8);
        assert_eq!(
            hex::encode(frames[1].header.hash),
            "ff584c07bbb0a039934a660bd49b7812af8ee847d1e675d9aba71c11fab3cfcb"
        );
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Defines the `IntoIter` owned iterator for arrays.

use crate::{
    fmt,
    iter::{self, ExactSizeIterator, FusedIterator, TrustedLen},
    mem::{self, MaybeUninit},
    ops::Range,
    ptr,
};

/// A by-value [array] iterator.
#[stable(feature = "array_value_iter", since = "1.51.0")]
pub struct IntoIter<T, const N: usize> {
    /// This is the array we are iterating over.
    ///
    /// Elements with index `i` where `alive.start <= i < alive.end` have not
    /// been yielded yet and are valid array entries. Elements with indices `i
    /// < alive.start` or `i >= alive.end` have been yielded already and must
    /// not be accessed anymore! Those dead elements might even be in a
    /// completely uninitialized state!
    ///
    /// So the invariants are:
    /// - `data[alive]` is alive (i.e. contains valid elements)
    /// - `data[..alive.start]` and `data[alive.end..]` are dead (i.e. the
    ///   elements were already read and must not be touched anymore!)
    data: [MaybeUninit<T>; N],

    /// The elements in `data` that have not been yielded yet.
    ///
    /// Invariants:
    /// - `alive.start <= alive.end`
    /// - `alive.end <= N`
    alive: Range<usize>,
}

impl<T, const N: usize> IntoIter<T, N> {
    /// Creates a new iterator over the given `array`.
    ///
    /// *Note*: this method might be deprecated in the future,
    /// after [`IntoIterator` is implemented for arrays][array-into-iter].
    ///
    /// # Examples
    ///
    /// ```
    /// use std::array;
    ///
    /// for value in array::IntoIter::new([1, 2, 3, 4, 5]) {
    ///     // The type of `value` is an `i32` here, instead of `&i32`
    ///     let _: i32 = value;
    /// }
    /// ```
    /// [array-into-iter]: https://github.com/rust-lang/rust/pull/65819
    #[stable(feature = "array_value_iter", since = "1.51.0")]
    pub fn new(array: [T; N]) -> Self {
        // SAFETY: The transmute here is actually safe. The docs of `MaybeUninit`
        // promise:
        //
        // > `MaybeUninit<T>` is guaranteed to have the same size and alignment
        // > as `T`.
        //
        // The docs even show a transmute from an array of `MaybeUninit<T>` to
        // an array of `T`.
        //
        // With that, this initialization satisfies the invariants.

        // FIXME(LukasKalbertodt): actually use `mem::transmute` here, once it
        // works with const generics:
        //     `mem::transmute::<[T; N], [MaybeUninit<T>; N]>(array)`
        //
        // Until then, we can use `mem::transmute_copy` to create a bitwise copy
        // as a different type, then forget `array` so that it is not dropped.
        unsafe {
            let iter = Self { data: mem::transmute_copy(&array), alive: 0..N };
            mem::forget(array);
            iter
        }
    }

    /// Returns an immutable slice of all elements that have not been yielded
    /// yet.
    #[stable(feature = "array_value_iter", since = "1.51.0")]
    pub fn as_slice(&self) -> &[T] {
        // SAFETY: We know that all elements within `alive` are properly initialized.
        unsafe {
            let slice = self.data.get_unchecked(self.alive.clone());
            MaybeUninit::slice_assume_init_ref(slice)
        }
    }

    /// Returns a mutable slice of all elements that have not been yielded yet.
    #[stable(feature = "array_value_iter", since = "1.51.0")]
    pub fn as_mut_slice(&mut self) -> &mut [T] {
        // SAFETY: We know that all elements within `alive` are properly initialized.
        unsafe {
            let slice = self.data.get_unchecked_mut(self.alive.clone());
            MaybeUninit::slice_assume_init_mut(slice)
        }
    }
}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T, const N: usize> Iterator for IntoIter<T, N> {
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        // Get the next index from the front.
        //
        // Increasing `alive.start` by 1 maintains the invariant regarding
        // `alive`. However, due to this change, for a short time, the alive
        // zone is not `data[alive]` anymore, but `data[idx..alive.end]`.
        self.alive.next().map(|idx| {
            // Read the element from the array.
            // SAFETY: `idx` is an index into the former "alive" region of the
            // array. Reading this element means that `data[idx]` is regarded as
            // dead now (i.e. do not touch). As `idx` was the start of the
            // alive-zone, the alive zone is now `data[alive]` again, restoring
            // all invariants.
            unsafe { self.data.get_unchecked(idx).assume_init_read() }
        })
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.len();
        (len, Some(len))
    }

    #[inline]
    fn fold<Acc, Fold>(mut self, init: Acc, mut fold: Fold) -> Acc
    where
        Fold: FnMut(Acc, Self::Item) -> Acc,
    {
        let data = &mut self.data;
        // FIXME: This uses try_fold(&mut iter) instead of fold(iter) because the latter
        //  would go through the blanket `impl Iterator for &mut I` implementation
        //  which lacks inline annotations on its methods and adding those would be a larger
        //  perturbation than using try_fold here.
        //  Whether it would be beneficial to add those annotations should be investigated separately.
        (&mut self.alive)
            .try_fold::<_, _, Result<_, !>>(init, |acc, idx| {
                // SAFETY: idx is obtained by folding over the `alive` range, which implies the
                // value is currently considered alive but as the range is being consumed each value
                // we read here will only be read once and then considered dead.
                Ok(fold(acc, unsafe { data.get_unchecked(idx).assume_init_read() }))
            })
            .unwrap()
    }

    fn count(self) -> usize {
        self.len()
    }

    fn last(mut self) -> Option<Self::Item> {
        self.next_back()
    }
}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
    fn next_back(&mut self) -> Option<Self::Item> {
        // Get the next index from the back.
        //
        // Decreasing `alive.end` by 1 maintains the invariant regarding
        // `alive`. However, due to this change, for a short time, the alive
        // zone is not `data[alive]` anymore, but `data[alive.start..=idx]`.
        self.alive.next_back().map(|idx| {
            // Read the element from the array.
            // SAFETY: `idx` is an index into the former "alive" region of the
            // array. Reading this element means that `data[idx]` is regarded as
            // dead now (i.e. do not touch). As `idx` was the end of the
            // alive-zone, the alive zone is now `data[alive]` again, restoring
            // all invariants.
            unsafe { self.data.get_unchecked(idx).assume_init_read() }
        })
    }
}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T, const N: usize> Drop for IntoIter<T, N> {
    fn drop(&mut self) {
        // SAFETY: This is safe: `as_mut_slice` returns exactly the sub-slice
        // of elements that have not been moved out yet and that remain
        // to be dropped.
        unsafe { ptr::drop_in_place(self.as_mut_slice()) }
    }
}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T, const N: usize> ExactSizeIterator for IntoIter<T, N> {
    fn len(&self) -> usize {
        // Will never underflow due to the invariant `alive.start <=
        // alive.end`.
        self.alive.end - self.alive.start
    }
    fn is_empty(&self) -> bool {
        self.alive.is_empty()
    }
}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T, const N: usize> FusedIterator for IntoIter<T, N> {}

// The iterator indeed reports the correct length. The number of "alive"
// elements (that will still be yielded) is the length of the range `alive`.
// This range is decremented in length in either `next` or `next_back`. It is
// always decremented by 1 in those methods, but only if `Some(_)` is returned.
#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
unsafe impl<T, const N: usize> TrustedLen for IntoIter<T, N> {}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T: Clone, const N: usize> Clone for IntoIter<T, N> {
    fn clone(&self) -> Self {
        // Note, we don't really need to match the exact same alive range, so
        // we can just clone into offset 0 regardless of where `self` is.
        let mut new = Self { data: MaybeUninit::uninit_array(), alive: 0..0 };

        // Clone all alive elements.
        for (src, dst) in iter::zip(self.as_slice(), &mut new.data) {
            // Write a clone into the new array, then update its alive range.
            // If cloning panics, we'll correctly drop the previous items.
            dst.write(src.clone());
            new.alive.end += 1;
        }

        new
    }
}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T: fmt::Debug, const N: usize> fmt::Debug for IntoIter<T, N> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Only print the elements that were not yielded yet: we cannot
        // access the yielded elements anymore.
        f.debug_tuple("IntoIter").field(&self.as_slice()).finish()
    }
}
<commit_msg>Update outdated docs of array::IntoIter::new.<commit_after>//! Defines the `IntoIter` owned iterator for arrays.

use crate::{
    fmt,
    iter::{self, ExactSizeIterator, FusedIterator, TrustedLen},
    mem::{self, MaybeUninit},
    ops::Range,
    ptr,
};

/// A by-value [array] iterator.
#[stable(feature = "array_value_iter", since = "1.51.0")]
pub struct IntoIter<T, const N: usize> {
    /// This is the array we are iterating over.
    ///
    /// Elements with index `i` where `alive.start <= i < alive.end` have not
    /// been yielded yet and are valid array entries. Elements with indices `i
    /// < alive.start` or `i >= alive.end` have been yielded already and must
    /// not be accessed anymore! Those dead elements might even be in a
    /// completely uninitialized state!
    ///
    /// So the invariants are:
    /// - `data[alive]` is alive (i.e. contains valid elements)
    /// - `data[..alive.start]` and `data[alive.end..]` are dead (i.e. the
    ///   elements were already read and must not be touched anymore!)
    data: [MaybeUninit<T>; N],

    /// The elements in `data` that have not been yielded yet.
    ///
    /// Invariants:
    /// - `alive.start <= alive.end`
    /// - `alive.end <= N`
    alive: Range<usize>,
}

impl<T, const N: usize> IntoIter<T, N> {
    /// Creates a new iterator over the given `array`.
    ///
    /// *Note*: this method might be deprecated in the future,
    /// since [`IntoIterator`] is now implemented for arrays.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::array;
    ///
    /// for value in array::IntoIter::new([1, 2, 3, 4, 5]) {
    ///     // The type of `value` is an `i32` here, instead of `&i32`
    ///     let _: i32 = value;
    /// }
    ///
    /// // Since Rust 1.53, arrays implement IntoIterator directly:
    /// for value in [1, 2, 3, 4, 5] {
    ///     // The type of `value` is an `i32` here, instead of `&i32`
    ///     let _: i32 = value;
    /// }
    /// ```
    #[stable(feature = "array_value_iter", since = "1.51.0")]
    pub fn new(array: [T; N]) -> Self {
        // SAFETY: The transmute here is actually safe. The docs of `MaybeUninit`
        // promise:
        //
        // > `MaybeUninit<T>` is guaranteed to have the same size and alignment
        // > as `T`.
        //
        // The docs even show a transmute from an array of `MaybeUninit<T>` to
        // an array of `T`.
        //
        // With that, this initialization satisfies the invariants.

        // FIXME(LukasKalbertodt): actually use `mem::transmute` here, once it
        // works with const generics:
        //     `mem::transmute::<[T; N], [MaybeUninit<T>; N]>(array)`
        //
        // Until then, we can use `mem::transmute_copy` to create a bitwise copy
        // as a different type, then forget `array` so that it is not dropped.
        unsafe {
            let iter = Self { data: mem::transmute_copy(&array), alive: 0..N };
            mem::forget(array);
            iter
        }
    }

    /// Returns an immutable slice of all elements that have not been yielded
    /// yet.
    #[stable(feature = "array_value_iter", since = "1.51.0")]
    pub fn as_slice(&self) -> &[T] {
        // SAFETY: We know that all elements within `alive` are properly initialized.
        unsafe {
            let slice = self.data.get_unchecked(self.alive.clone());
            MaybeUninit::slice_assume_init_ref(slice)
        }
    }

    /// Returns a mutable slice of all elements that have not been yielded yet.
    #[stable(feature = "array_value_iter", since = "1.51.0")]
    pub fn as_mut_slice(&mut self) -> &mut [T] {
        // SAFETY: We know that all elements within `alive` are properly initialized.
        unsafe {
            let slice = self.data.get_unchecked_mut(self.alive.clone());
            MaybeUninit::slice_assume_init_mut(slice)
        }
    }
}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T, const N: usize> Iterator for IntoIter<T, N> {
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        // Get the next index from the front.
        //
        // Increasing `alive.start` by 1 maintains the invariant regarding
        // `alive`. However, due to this change, for a short time, the alive
        // zone is not `data[alive]` anymore, but `data[idx..alive.end]`.
        self.alive.next().map(|idx| {
            // Read the element from the array.
            // SAFETY: `idx` is an index into the former "alive" region of the
            // array. Reading this element means that `data[idx]` is regarded as
            // dead now (i.e. do not touch). As `idx` was the start of the
            // alive-zone, the alive zone is now `data[alive]` again, restoring
            // all invariants.
            unsafe { self.data.get_unchecked(idx).assume_init_read() }
        })
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.len();
        (len, Some(len))
    }

    #[inline]
    fn fold<Acc, Fold>(mut self, init: Acc, mut fold: Fold) -> Acc
    where
        Fold: FnMut(Acc, Self::Item) -> Acc,
    {
        let data = &mut self.data;
        // FIXME: This uses try_fold(&mut iter) instead of fold(iter) because the latter
        //  would go through the blanket `impl Iterator for &mut I` implementation
        //  which lacks inline annotations on its methods and adding those would be a larger
        //  perturbation than using try_fold here.
        //  Whether it would be beneficial to add those annotations should be investigated separately.
        (&mut self.alive)
            .try_fold::<_, _, Result<_, !>>(init, |acc, idx| {
                // SAFETY: idx is obtained by folding over the `alive` range, which implies the
                // value is currently considered alive but as the range is being consumed each value
                // we read here will only be read once and then considered dead.
                Ok(fold(acc, unsafe { data.get_unchecked(idx).assume_init_read() }))
            })
            .unwrap()
    }

    fn count(self) -> usize {
        self.len()
    }

    fn last(mut self) -> Option<Self::Item> {
        self.next_back()
    }
}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
    fn next_back(&mut self) -> Option<Self::Item> {
        // Get the next index from the back.
        //
        // Decreasing `alive.end` by 1 maintains the invariant regarding
        // `alive`. However, due to this change, for a short time, the alive
        // zone is not `data[alive]` anymore, but `data[alive.start..=idx]`.
        self.alive.next_back().map(|idx| {
            // Read the element from the array.
            // SAFETY: `idx` is an index into the former "alive" region of the
            // array. Reading this element means that `data[idx]` is regarded as
            // dead now (i.e. do not touch). As `idx` was the end of the
            // alive-zone, the alive zone is now `data[alive]` again, restoring
            // all invariants.
            unsafe { self.data.get_unchecked(idx).assume_init_read() }
        })
    }
}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T, const N: usize> Drop for IntoIter<T, N> {
    fn drop(&mut self) {
        // SAFETY: This is safe: `as_mut_slice` returns exactly the sub-slice
        // of elements that have not been moved out yet and that remain
        // to be dropped.
        unsafe { ptr::drop_in_place(self.as_mut_slice()) }
    }
}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T, const N: usize> ExactSizeIterator for IntoIter<T, N> {
    fn len(&self) -> usize {
        // Will never underflow due to the invariant `alive.start <=
        // alive.end`.
        self.alive.end - self.alive.start
    }
    fn is_empty(&self) -> bool {
        self.alive.is_empty()
    }
}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T, const N: usize> FusedIterator for IntoIter<T, N> {}

// The iterator indeed reports the correct length. The number of "alive"
// elements (that will still be yielded) is the length of the range `alive`.
// This range is decremented in length in either `next` or `next_back`. It is
// always decremented by 1 in those methods, but only if `Some(_)` is returned.
#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
unsafe impl<T, const N: usize> TrustedLen for IntoIter<T, N> {}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T: Clone, const N: usize> Clone for IntoIter<T, N> {
    fn clone(&self) -> Self {
        // Note, we don't really need to match the exact same alive range, so
        // we can just clone into offset 0 regardless of where `self` is.
        let mut new = Self { data: MaybeUninit::uninit_array(), alive: 0..0 };

        // Clone all alive elements.
        for (src, dst) in iter::zip(self.as_slice(), &mut new.data) {
            // Write a clone into the new array, then update its alive range.
            // If cloning panics, we'll correctly drop the previous items.
            dst.write(src.clone());
            new.alive.end += 1;
        }

        new
    }
}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T: fmt::Debug, const N: usize> fmt::Debug for IntoIter<T, N> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Only print the elements that were not yielded yet: we cannot
        // access the yielded elements anymore.
        f.debug_tuple("IntoIter").field(&self.as_slice()).finish()
    }
}
<|endoftext|>