<commit_before>// Copyright (c) 2016-2020 Fabian Schuiki

use heck::CamelCase;
use proc_macro::TokenStream;
use quote::{format_ident, quote, ToTokens};
use std::{
    cell::RefCell,
    collections::{BTreeSet, HashSet},
};
use syn::visit::Visit;

// CAUTION: This is all wildly unstable and relies on the compiler maintaining
// a certain order between proc macro expansions. So this could break any
// minute. Better have a robust CI.
thread_local! {
    static QUERIES: RefCell<Vec<String>> = Default::default();
}

pub(crate) fn mark_query(_args: TokenStream, input: TokenStream) -> TokenStream {
    // Parse the input.
    let input = syn::parse_macro_input!(input as syn::ItemFn);

    // Map everything to a string here. Compiler panics horribly if we hand out
    // the actual idents and generics.
    QUERIES.with(|c| c.borrow_mut().push(input.to_token_stream().to_string()));

    // Produce some output.
    let note = format!("call query via database: cx.{}(...)", input.sig.ident);
    let output = quote! {
        #[deprecated(note = #note)]
        #input
    };
    output.into()
}

pub(crate) fn derive_query_db(input: TokenStream) -> TokenStream {
    let input = proc_macro2::TokenStream::from(input);
    let mut output = proc_macro2::TokenStream::new();

    // Flush the accumulated queries.
    let queries = QUERIES.with(|c| std::mem::replace(&mut *c.borrow_mut(), Default::default()));

    // Make a set of conflicting argument names.
    let reserved: HashSet<_> = ["query_key", "query_storage", "query_tag", "query_result"]
        .iter()
        .copied()
        .collect();

    // Collect all query lifetimes.
    let mut ltset = BTreeSet::new();

    // Process the queries.
    let mut funcs = vec![];
    let mut caches = vec![];
    let mut tags = vec![];
    let mut tag_debugs = vec![];
    let mut keys = vec![];

    for raw_query in &queries {
        // Parse the fn.
        let query: syn::ItemFn = syn::parse_str(raw_query).unwrap();

        // Disect a few things.
        let name = query.sig.ident.clone();
        let mut generics = query.sig.generics.clone();
        let args = query.sig.inputs.iter().skip(1);
        let result = match query.sig.output {
            syn::ReturnType::Type(_, ty) => ty.as_ref().clone(),
            _ => panic!("query {} has no return type", name),
        };

        // Filter out the doc comments such that we can apply them to the trait
        // fn as well.
        let doc_attrs = query.attrs.iter().filter(|a| a.path.is_ident("doc"));

        // Determine the argument list of the query.
        let arg_pats = args.clone().map(|arg| match arg {
            syn::FnArg::Typed(pat) => pat,
            _ => unreachable!("there should be no self args left at this point"),
        });
        let arg_names: Vec<_> = arg_pats
            .clone()
            .enumerate()
            .map(|(i, pat)| match pat.pat.as_ref() {
                syn::Pat::Ident(id) if !reserved.contains(id.ident.to_string().as_str()) => {
                    id.ident.clone()
                }
                _ => format_ident!("arg{}", i),
            })
            .collect();
        let arg_types: Vec<_> = arg_pats
            .clone()
            .map(|pat| pat.ty.as_ref().clone())
            .collect();

        // Collect lifetimes and strip them from the generics.
        let original_generics = generics.clone();
        for param in std::mem::replace(&mut generics.params, Default::default()) {
            match param {
                syn::GenericParam::Lifetime(ltdef) => {
                    ltset.insert(ltdef.lifetime);
                }
                _ => generics.params.push(param),
            }
        }

        // Determine the generics that go into the key.
        struct KeyGenerics(HashSet<String>);
        impl Visit<'_> for KeyGenerics {
            fn visit_lifetime(&mut self, arg: &syn::Lifetime) {
                self.0.insert(arg.to_string());
            }

            fn visit_generic_argument(&mut self, arg: &syn::GenericArgument) {
                match arg {
                    syn::GenericArgument::Lifetime(x) => {
                        self.0.insert(x.to_string());
                    }
                    syn::GenericArgument::Type(syn::Type::Path(x)) => {
                        if let Some(x) = x.path.get_ident() {
                            self.0.insert(x.to_string());
                        }
                    }
                    _ => (),
                }
            }
        }

        let mut key_generics_set = KeyGenerics(Default::default());
        for arg_type in &arg_types {
            key_generics_set.visit_type(&arg_type);
        }
        // println!("Key generics of {}: {:?}", name, key_generics_set.0);

        let mut key_generics = syn::Generics {
            lt_token: original_generics.lt_token,
            params: Default::default(),
            gt_token: original_generics.gt_token,
            where_clause: None,
        };
        for param in original_generics.params.iter() {
            let keep = match param {
                syn::GenericParam::Type(x) => key_generics_set.0.contains(&x.ident.to_string()),
                syn::GenericParam::Lifetime(x) => {
                    key_generics_set.0.contains(&x.lifetime.to_string())
                }
                syn::GenericParam::Const(x) => key_generics_set.0.contains(&x.ident.to_string()),
            };
            if keep {
                key_generics.params.push(param.clone());
            }
        }
        // println!("Key generics: {}", quote! { #key_generics });

        // Render the key type that we use to look up things in the database.
        let key_name = format_ident!("{}QueryKey", name.to_string().to_camel_case());
        let key_indices = (0..arg_types.len()).map(syn::Index::from);
        let doc = format!("The arguments passed to the `{}` query.", name);
        keys.push(quote! {
            #[doc = #doc]
            #[derive(Clone, PartialEq, Eq, Hash)]
            pub struct #key_name #key_generics (#(#arg_types),*);

            impl #key_generics std::fmt::Debug for #key_name #key_generics {
                fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
                    f.debug_tuple("")#(.field(&self.#key_indices))*.finish()
                }
            }
        });

        // Determine the cache field name.
        let cache_name = format_ident!("cached_{}", name);

        // Render a query tag that can be pushed onto the query stack to break
        // cycles.
        let tag_name = format_ident!("{}", name.to_string().to_camel_case());
        let doc = format!("The `{}` query.", name);
        tags.push(quote! {
            #[doc = #doc]
            #tag_name (#key_name #key_generics),
        });
        tag_debugs.push(quote! {
            QueryTag::#tag_name (x) => write!(f, "{}{:?}", stringify!(#name), x),
        });

        // Render the query for the database trait.
        funcs.push(quote! {
            #(#doc_attrs)*
            fn #name #generics (&self, #(#arg_names: #arg_types),*) -> #result {
                let query_storage = self.storage();
                let query_key = #key_name(#(Clone::clone(&#arg_names)),*);
                let query_tag = QueryTag::#tag_name(query_key.clone());
                self.before_query(&query_tag);

                // Check if we already have a result for this query.
                if let Some(result) = query_storage.#cache_name.borrow().get(&query_key) {
                    trace!("Serving {}{:?} from cache", stringify!(#name), query_key);
                    return Clone::clone(&result);
                }
                trace!("Executing {}{:?}", stringify!(#name), query_key);

                // Push the query onto the stack, checking if it is already in
                // flight.
                query_storage.stack.borrow_mut().push(query_tag.clone());
                if !query_storage.inflight.borrow_mut().insert(query_tag.clone()) {
                    self.handle_cycle();
                    // The above never returns.
                }

                // Execute the query.
                #[allow(deprecated)]
                let query_result = #name(self.context(), #(#arg_names),*);
                query_storage.#cache_name.borrow_mut().insert(query_key, Clone::clone(&query_result));

                // Pop the query from the stack.
                query_storage.inflight.borrow_mut().remove(&query_tag);
                query_storage.stack.borrow_mut().pop();
                self.after_query(&query_tag);
                query_result
            }
        });

        // Render the cache field for the storage struct.
        let doc = format!("Cached results of the `{}` query.", name);
        caches.push(quote! {
            #[doc = #doc]
            pub #cache_name: RefCell<HashMap<#key_name #key_generics, #result>>,
        });
    }

    // Extract query lifetimes.
    let lts = if ltset.is_empty() {
        quote! {}
    } else {
        let lts = ltset.iter();
        quote! { <#(#lts),*> }
    };

    // Generate the query database trait.
    output.extend(quote! {
        #input
        pub trait QueryDatabase #lts {
            /// The type passed as first argument to compiler query
            /// implementations.
            type Context: crate::Context #lts;

            /// Get the context that is passed to compiler query
            /// implementations.
            fn context(&self) -> &Self::Context;

            /// Get the query caches and runtime data.
            fn storage(&self) -> &QueryStorage #lts;

            /// Called when a query cycle is detected.
            fn handle_cycle(&self) -> ! {
                error!("Query cycle detected:");
                for x in self.storage().stack.borrow().iter() {
                    error!(" - {:?}", x);
                }
                panic!("query cycle detected");
            }

            /// Called before a query is executed.
            fn before_query(&self, tag: &QueryTag #lts) {}

            /// Called after a query is executed.
            fn after_query(&self, tag: &QueryTag #lts) {}

            #(#funcs)*
        }
    });

    // Generate the query storage struct.
    output.extend(quote! {
        /// A collection of query caches and runtime data for a `QueryDatabase`.
        #[derive(Default)]
        pub struct QueryStorage #lts {
            /// A stack of the currently-executing queries.
            pub stack: RefCell<Vec<QueryTag #lts>>,
            /// A set of the currently-executing queries.
            pub inflight: RefCell<HashSet<QueryTag #lts>>,

            #(#caches)*
        }
    });

    // Generate the query tag enum.
    output.extend(quote! {
        /// A tag identifying any of the queries in `QueryDatabase`.
        #[derive(Clone, Hash, Eq, PartialEq)]
        pub enum QueryTag #lts {
            #(#tags)*
        }

        impl #lts std::fmt::Debug for QueryTag #lts {
            fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
                match self {
                    #(#tag_debugs)*
                }
            }
        }
    });

    // Generate the query keys.
    output.extend(quote! {
        #(#keys)*
    });

    // Produce some output.
    // println!("{}", output);
    output.into()
}
<commit_msg>derive: Fix lifetimes added to query keys<commit_after>// Copyright (c) 2016-2020 Fabian Schuiki

use heck::CamelCase;
use proc_macro::TokenStream;
use quote::{format_ident, quote, ToTokens};
use std::{
    cell::RefCell,
    collections::{BTreeSet, HashSet},
};
use syn::visit::Visit;

// CAUTION: This is all wildly unstable and relies on the compiler maintaining
// a certain order between proc macro expansions. So this could break any
// minute. Better have a robust CI.
thread_local! {
    static QUERIES: RefCell<Vec<String>> = Default::default();
}

pub(crate) fn mark_query(_args: TokenStream, input: TokenStream) -> TokenStream {
    // Parse the input.
    let input = syn::parse_macro_input!(input as syn::ItemFn);

    // Map everything to a string here. Compiler panics horribly if we hand out
    // the actual idents and generics.
    QUERIES.with(|c| c.borrow_mut().push(input.to_token_stream().to_string()));

    // Produce some output.
    let note = format!("call query via database: cx.{}(...)", input.sig.ident);
    let output = quote! {
        #[deprecated(note = #note)]
        #input
    };
    output.into()
}

pub(crate) fn derive_query_db(input: TokenStream) -> TokenStream {
    let input = proc_macro2::TokenStream::from(input);
    let mut output = proc_macro2::TokenStream::new();

    // Flush the accumulated queries.
    let queries = QUERIES.with(|c| std::mem::replace(&mut *c.borrow_mut(), Default::default()));

    // Make a set of conflicting argument names.
    let reserved: HashSet<_> = ["query_key", "query_storage", "query_tag", "query_result"]
        .iter()
        .copied()
        .collect();

    // Collect all query lifetimes.
    let mut ltset = BTreeSet::new();

    // Process the queries.
    let mut funcs = vec![];
    let mut caches = vec![];
    let mut tags = vec![];
    let mut tag_debugs = vec![];
    let mut keys = vec![];

    for raw_query in &queries {
        // Parse the fn.
        let query: syn::ItemFn = syn::parse_str(raw_query).unwrap();

        // Disect a few things.
        let name = query.sig.ident.clone();
        let mut generics = query.sig.generics.clone();
        let args = query.sig.inputs.iter().skip(1);
        let result = match query.sig.output {
            syn::ReturnType::Type(_, ty) => ty.as_ref().clone(),
            _ => panic!("query {} has no return type", name),
        };

        // Filter out the doc comments such that we can apply them to the trait
        // fn as well.
        let doc_attrs = query.attrs.iter().filter(|a| a.path.is_ident("doc"));

        // Determine the argument list of the query.
        let arg_pats = args.clone().map(|arg| match arg {
            syn::FnArg::Typed(pat) => pat,
            _ => unreachable!("there should be no self args left at this point"),
        });
        let arg_names: Vec<_> = arg_pats
            .clone()
            .enumerate()
            .map(|(i, pat)| match pat.pat.as_ref() {
                syn::Pat::Ident(id) if !reserved.contains(id.ident.to_string().as_str()) => {
                    id.ident.clone()
                }
                _ => format_ident!("arg{}", i),
            })
            .collect();
        let arg_types: Vec<_> = arg_pats
            .clone()
            .map(|pat| pat.ty.as_ref().clone())
            .collect();

        // Collect lifetimes and strip them from the generics.
        let original_generics = generics.clone();
        for param in std::mem::replace(&mut generics.params, Default::default()) {
            match param {
                syn::GenericParam::Lifetime(ltdef) => {
                    ltset.insert(ltdef.lifetime);
                }
                _ => generics.params.push(param),
            }
        }

        // Determine the generics that go into the key.
        struct KeyGenerics(HashSet<String>);
        impl Visit<'_> for KeyGenerics {
            fn visit_lifetime(&mut self, node: &syn::Lifetime) {
                self.0.insert(node.to_string());
                syn::visit::visit_lifetime(self, node);
            }

            fn visit_type(&mut self, node: &syn::Type) {
                match node {
                    syn::Type::Path(x) => {
                        if let Some(x) = x.path.get_ident() {
                            self.0.insert(x.to_string());
                        }
                    }
                    _ => (),
                }
                syn::visit::visit_type(self, node);
            }
        }

        let mut key_generics_set = KeyGenerics(Default::default());
        for arg_type in &arg_types {
            key_generics_set.visit_type(&arg_type);
        }
        // println!("Key generics of {}: {:?}", name, key_generics_set.0);

        let mut key_generics = syn::Generics {
            lt_token: original_generics.lt_token,
            params: Default::default(),
            gt_token: original_generics.gt_token,
            where_clause: None,
        };
        for param in original_generics.params.iter() {
            let keep = match param {
                syn::GenericParam::Type(x) => key_generics_set.0.contains(&x.ident.to_string()),
                syn::GenericParam::Lifetime(x) => {
                    key_generics_set.0.contains(&x.lifetime.to_string())
                }
                syn::GenericParam::Const(x) => key_generics_set.0.contains(&x.ident.to_string()),
            };
            if keep {
                key_generics.params.push(param.clone());
            }
        }
        // println!("Key generics: {}", quote! { #key_generics });

        // Render the key type that we use to look up things in the database.
        let key_name = format_ident!("{}QueryKey", name.to_string().to_camel_case());
        let key_indices = (0..arg_types.len()).map(syn::Index::from);
        let doc = format!("The arguments passed to the `{}` query.", name);
        keys.push(quote! {
            #[doc = #doc]
            #[derive(Clone, PartialEq, Eq, Hash)]
            pub struct #key_name #key_generics (#(#arg_types),*);

            impl #key_generics std::fmt::Debug for #key_name #key_generics {
                fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
                    f.debug_tuple("")#(.field(&self.#key_indices))*.finish()
                }
            }
        });

        // Determine the cache field name.
        let cache_name = format_ident!("cached_{}", name);

        // Render a query tag that can be pushed onto the query stack to break
        // cycles.
        let tag_name = format_ident!("{}", name.to_string().to_camel_case());
        let doc = format!("The `{}` query.", name);
        tags.push(quote! {
            #[doc = #doc]
            #tag_name (#key_name #key_generics),
        });
        tag_debugs.push(quote! {
            QueryTag::#tag_name (x) => write!(f, "{}{:?}", stringify!(#name), x),
        });

        // Render the query for the database trait.
        funcs.push(quote! {
            #(#doc_attrs)*
            fn #name #generics (&self, #(#arg_names: #arg_types),*) -> #result {
                let query_storage = self.storage();
                let query_key = #key_name(#(Clone::clone(&#arg_names)),*);
                let query_tag = QueryTag::#tag_name(query_key.clone());
                self.before_query(&query_tag);

                // Check if we already have a result for this query.
                if let Some(result) = query_storage.#cache_name.borrow().get(&query_key) {
                    trace!("Serving {}{:?} from cache", stringify!(#name), query_key);
                    return Clone::clone(&result);
                }
                trace!("Executing {}{:?}", stringify!(#name), query_key);

                // Push the query onto the stack, checking if it is already in
                // flight.
                query_storage.stack.borrow_mut().push(query_tag.clone());
                if !query_storage.inflight.borrow_mut().insert(query_tag.clone()) {
                    self.handle_cycle();
                    // The above never returns.
                }

                // Execute the query.
                #[allow(deprecated)]
                let query_result = #name(self.context(), #(#arg_names),*);
                query_storage.#cache_name.borrow_mut().insert(query_key, Clone::clone(&query_result));

                // Pop the query from the stack.
                query_storage.inflight.borrow_mut().remove(&query_tag);
                query_storage.stack.borrow_mut().pop();
                self.after_query(&query_tag);
                query_result
            }
        });

        // Render the cache field for the storage struct.
        let doc = format!("Cached results of the `{}` query.", name);
        caches.push(quote! {
            #[doc = #doc]
            pub #cache_name: RefCell<HashMap<#key_name #key_generics, #result>>,
        });
    }

    // Extract query lifetimes.
    let lts = if ltset.is_empty() {
        quote! {}
    } else {
        let lts = ltset.iter();
        quote! { <#(#lts),*> }
    };

    // Generate the query database trait.
    output.extend(quote! {
        #input
        pub trait QueryDatabase #lts {
            /// The type passed as first argument to compiler query
            /// implementations.
            type Context: crate::Context #lts;

            /// Get the context that is passed to compiler query
            /// implementations.
            fn context(&self) -> &Self::Context;

            /// Get the query caches and runtime data.
            fn storage(&self) -> &QueryStorage #lts;

            /// Called when a query cycle is detected.
            fn handle_cycle(&self) -> ! {
                error!("Query cycle detected:");
                for x in self.storage().stack.borrow().iter() {
                    error!(" - {:?}", x);
                }
                panic!("query cycle detected");
            }

            /// Called before a query is executed.
            fn before_query(&self, tag: &QueryTag #lts) {}

            /// Called after a query is executed.
            fn after_query(&self, tag: &QueryTag #lts) {}

            #(#funcs)*
        }
    });

    // Generate the query storage struct.
    output.extend(quote! {
        /// A collection of query caches and runtime data for a `QueryDatabase`.
        #[derive(Default)]
        pub struct QueryStorage #lts {
            /// A stack of the currently-executing queries.
            pub stack: RefCell<Vec<QueryTag #lts>>,
            /// A set of the currently-executing queries.
            pub inflight: RefCell<HashSet<QueryTag #lts>>,

            #(#caches)*
        }
    });

    // Generate the query tag enum.
    output.extend(quote! {
        /// A tag identifying any of the queries in `QueryDatabase`.
        #[derive(Clone, Hash, Eq, PartialEq)]
        pub enum QueryTag #lts {
            #(#tags)*
        }

        impl #lts std::fmt::Debug for QueryTag #lts {
            fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
                match self {
                    #(#tag_debugs)*
                }
            }
        }
    });

    // Generate the query keys.
    output.extend(quote! {
        #(#keys)*
    });

    // Produce some output.
    // println!("{}", output);
    output.into()
}
<|endoftext|>


<commit_before>//! A Rust interface for Objective-C blocks.
//!
//! For more information on the specifics of the block implementation, see
//! Clang's documentation: http://clang.llvm.org/docs/Block-ABI-Apple.html

use std::mem;
use std::ptr;
use libc::{c_int, c_ulong};

use runtime::Class;
use {Id, Message};

#[allow(improper_ctypes)]
#[link(name = "Foundation", kind = "framework")]
extern {
    static _NSConcreteStackBlock: Class;
}

/// An invoke function for a `Block`; this is the raw C function called by the
/// Objective-C runtime.
type BlockInvoke<A, R> = unsafe extern fn(*mut Block<A, R>, ...) -> R;

/// An invoke function for a `ConcreteBlock`; this is the raw C function called
/// by the Objective-C runtime.
pub type ConcreteBlockInvoke<A, R, C> =
    unsafe extern fn(*mut ConcreteBlock<A, R, C>, ...) -> R;

/// Types that may be used as the arguments to an Objective-C block.
pub trait BlockArguments {
    /// Calls the given `Block` with self as the arguments.
    fn call_block<R>(self, block: &Block<Self, R>) -> R;

    /// Returns an invoke function for a `ConcreteBlock` that takes this type
    /// of arguments.
    fn invoke_for_concrete_block<R, C: Clone>() -> ConcreteBlockInvoke<Self, R, C>;
}

macro_rules! block_args_impl(
    ($f:ident $(, $a:ident : $t:ident)*) => (
        impl<$($t),*> BlockArguments for ($($t,)*) {
            fn call_block<R>(self, block: &Block<($($t,)*), R>) -> R {
                let invoke: unsafe extern fn(*mut Block<($($t,)*), R> $(, $t)*) -> R = unsafe {
                    mem::transmute(block.invoke)
                };
                let ($($a,)*) = self;
                let block_ptr = block as *const _ as *mut _;
                unsafe {
                    invoke(block_ptr $(, $a)*)
                }
            }

            fn invoke_for_concrete_block<R, X: Clone>() ->
                    ConcreteBlockInvoke<($($t,)*), R, X> {
                unsafe extern fn $f<R, X: Clone $(, $t)*>(
                        block_ptr: *mut ConcreteBlock<($($t,)*), R, X>
                        $(, $a: $t)*) -> R {
                    let args = ($($a,)*);
                    let block = &*block_ptr;
                    (block.rust_invoke)(&block.context, args)
                }

                unsafe {
                    mem::transmute($f::<R, X $(, $t)*>)
                }
            }
        }
    );
);

block_args_impl!(concrete_block_invoke_args0);
block_args_impl!(concrete_block_invoke_args1, a: A);
block_args_impl!(concrete_block_invoke_args2, a: A, b: B);
block_args_impl!(concrete_block_invoke_args3, a: A, b: B, c: C);
block_args_impl!(concrete_block_invoke_args4, a: A, b: B, c: C, d: D);
block_args_impl!(concrete_block_invoke_args5, a: A, b: B, c: C, d: D, e: E);
block_args_impl!(concrete_block_invoke_args6, a: A, b: B, c: C, d: D, e: E, f: F);
block_args_impl!(concrete_block_invoke_args7, a: A, b: B, c: C, d: D, e: E, f: F, g: G);
block_args_impl!(concrete_block_invoke_args8, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H);
block_args_impl!(concrete_block_invoke_args9, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I);
block_args_impl!(concrete_block_invoke_args10, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J);
block_args_impl!(concrete_block_invoke_args11, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K);
block_args_impl!(concrete_block_invoke_args12, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K, l: L);

/// An Objective-C block that takes arguments of `A` when called and
/// returns a value of `R`.
#[repr(C)]
pub struct Block<A: BlockArguments, R> {
    isa: *const Class,
    flags: c_int,
    _reserved: c_int,
    invoke: BlockInvoke<A, R>,
}

impl<A: BlockArguments, R> Block<A, R> {
    /// Copy self onto the heap.
    pub fn copy(&self) -> Id<Block<A, R>> {
        unsafe {
            let block = msg_send![self copy] as *mut Block<A, R>;
            Id::from_retained_ptr(block)
        }
    }

    /// Call self with the given arguments.
    pub fn call(&self, args: A) -> R {
        args.call_block(self)
    }
}

impl<A: BlockArguments, R> Message for Block<A, R> { }

/// An Objective-C block whose size is known at compile time and may be
/// constructed on the stack.
#[repr(C)]
pub struct ConcreteBlock<A: BlockArguments, R, C: Clone> {
    base: Block<A, R>,
    descriptor: *const BlockDescriptor<ConcreteBlock<A, R, C>>,
    rust_invoke: fn (&C, A) -> R,
    context: C,
}

impl<A: BlockArguments, R, C: Clone> ConcreteBlock<A, R, C> {
    /// Constructs a `ConcreteBlock` with the given invoke function and context.
    /// When the block is called, it will return the value that results from
    /// calling the invoke function with a reference to its context.
    pub fn new(invoke: fn (&C, A) -> R, context: C) -> ConcreteBlock<A, R, C> {
        let extern_invoke: ConcreteBlockInvoke<A, R, C> =
            BlockArguments::invoke_for_concrete_block();
        ConcreteBlock {
            base: Block {
                isa: &_NSConcreteStackBlock,
                // 1 << 25 = BLOCK_HAS_COPY_DISPOSE
                flags: 1 << 25,
                _reserved: 0,
                invoke: unsafe { mem::transmute(extern_invoke) },
            },
            // TODO: don't leak memory here
            descriptor: unsafe {
                mem::transmute(box BlockDescriptor::<A, R, C>::new())
            },
            rust_invoke: invoke,
            context: context,
        }
    }
}

impl<A: BlockArguments, R, C: Clone> Clone for ConcreteBlock<A, R, C> {
    fn clone(&self) -> ConcreteBlock<A, R, C> {
        ConcreteBlock::new(self.rust_invoke, self.context.clone())
    }
}

impl<A: BlockArguments, R, C: Clone> Deref<Block<A, R>> for ConcreteBlock<A, R, C> {
    fn deref(&self) -> &Block<A, R> {
        &self.base
    }
}

unsafe extern fn block_context_dispose<A: BlockArguments, R, C: Clone>(
        block: &mut ConcreteBlock<A, R, C>) {
    // Read the context from the block and let it drop
    ptr::read(&block.context);
}

unsafe extern fn block_context_copy<A: BlockArguments, R, C: Clone>(
        dst: &mut ConcreteBlock<A, R, C>, src: &ConcreteBlock<A, R, C>) {
    // The src block actually gets memmoved to the destination beforehand,
    // but we'll set the function pointer, too, to be safe.
    ptr::write(&mut dst.rust_invoke, src.rust_invoke);
    ptr::write(&mut dst.context, src.context.clone());
}

#[repr(C)]
struct BlockDescriptor<B> {
    _reserved: c_ulong,
    block_size: c_ulong,
    copy_helper: unsafe extern fn(&mut B, &B),
    dispose_helper: unsafe extern fn(&mut B),
}

impl<A: BlockArguments, R, C: Clone> BlockDescriptor<ConcreteBlock<A, R, C>> {
    fn new() -> BlockDescriptor<ConcreteBlock<A, R, C>> {
        BlockDescriptor {
            _reserved: 0,
            block_size: mem::size_of::<ConcreteBlock<A, R, C>>() as c_ulong,
            copy_helper: block_context_copy::<A, R, C>,
            dispose_helper: block_context_dispose::<A, R, C>,
        }
    }
}

#[cfg(test)]
mod tests {
    use Id;
    use objc_test_utils;
    use super::{Block, ConcreteBlock};

    fn get_int_block() -> &'static Block<(), int> {
        unsafe {
            &*(objc_test_utils::get_int_block() as *const _)
        }
    }

    fn get_int_block_with(i: int) -> Id<Block<(), int>> {
        unsafe {
            let ptr = objc_test_utils::get_int_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn get_add_block() -> &'static Block<(int,), int> {
        unsafe {
            &*(objc_test_utils::get_add_block() as *const _)
        }
    }

    fn get_add_block_with(i: int) -> Id<Block<(int,), int>> {
        unsafe {
            let ptr = objc_test_utils::get_add_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn invoke_int_block(block: &Block<(), int>) -> int {
        let ptr = block as *const _ as *const _;
        unsafe {
            objc_test_utils::invoke_int_block(ptr)
        }
    }

    fn invoke_add_block(block: &Block<(int,), int>, a: int) -> int {
        let ptr = block as *const _ as *const _;
        unsafe {
            objc_test_utils::invoke_add_block(ptr, a)
        }
    }

    #[test]
    fn test_call_block() {
        let block = get_int_block();
        assert!(block.call(()) == 7);

        let block = get_int_block_with(13);
        assert!(block.call(()) == 13);
    }

    #[test]
    fn test_call_block_args() {
        let block = get_add_block();
        assert!(block.call((2,)) == 9);

        let block = get_add_block_with(13);
        assert!(block.call((2,)) == 15);
    }

    #[test]
    fn test_create_block() {
        fn block_get_int(context: &int, _args: ()) -> int {
            *context
        }

        let block = ConcreteBlock::new(block_get_int, 13);
        let result = invoke_int_block(&*block);
        assert!(result == 13);
    }

    #[test]
    fn test_create_block_args() {
        fn block_add_int(context: &int, (a,): (int,)) -> int {
            a + *context
        }

        let block = ConcreteBlock::new(block_add_int, 5);
        let result = invoke_add_block(&*block, 6);
        assert!(result == 11);
    }

    #[test]
    fn test_concrete_block_clone() {
        fn block_get_string_len(context: &String, _args: ()) -> uint {
            context.len()
        }

        let s = "Hello!".into_string();
        let expected_len = s.len();
        let block = ConcreteBlock::new(block_get_string_len, s);
        assert!(block.call(()) == expected_len);

        let cloned = block.clone();
        assert!(cloned.call(()) == expected_len);

        drop(block);
        assert!(cloned.call(()) == expected_len);
    }

    #[test]
    fn test_concrete_block_copy() {
        fn block_get_string_len(context: &String, _args: ()) -> uint {
            context.len()
        }

        let s = "Hello!".into_string();
        let expected_len = s.len();
        let block = ConcreteBlock::new(block_get_string_len, s);
        assert!(block.call(()) == expected_len);

        let copied = block.copy();
        assert!(copied.call(()) == expected_len);

        drop(block);
        assert!(copied.call(()) == expected_len);
    }
}
<commit_msg>Rearranged some test code.<commit_after>//! A Rust interface for Objective-C blocks.
//!
//! For more information on the specifics of the block implementation, see
//! Clang's documentation: http://clang.llvm.org/docs/Block-ABI-Apple.html

use std::mem;
use std::ptr;
use libc::{c_int, c_ulong};

use runtime::Class;
use {Id, Message};

#[allow(improper_ctypes)]
#[link(name = "Foundation", kind = "framework")]
extern {
    static _NSConcreteStackBlock: Class;
}

/// An invoke function for a `Block`; this is the raw C function called by the
/// Objective-C runtime.
type BlockInvoke<A, R> = unsafe extern fn(*mut Block<A, R>, ...) -> R;

/// An invoke function for a `ConcreteBlock`; this is the raw C function called
/// by the Objective-C runtime.
pub type ConcreteBlockInvoke<A, R, C> =
    unsafe extern fn(*mut ConcreteBlock<A, R, C>, ...) -> R;

/// Types that may be used as the arguments to an Objective-C block.
pub trait BlockArguments {
    /// Calls the given `Block` with self as the arguments.
    fn call_block<R>(self, block: &Block<Self, R>) -> R;

    /// Returns an invoke function for a `ConcreteBlock` that takes this type
    /// of arguments.
    fn invoke_for_concrete_block<R, C: Clone>() -> ConcreteBlockInvoke<Self, R, C>;
}

macro_rules! block_args_impl(
    ($f:ident $(, $a:ident : $t:ident)*) => (
        impl<$($t),*> BlockArguments for ($($t,)*) {
            fn call_block<R>(self, block: &Block<($($t,)*), R>) -> R {
                let invoke: unsafe extern fn(*mut Block<($($t,)*), R> $(, $t)*) -> R = unsafe {
                    mem::transmute(block.invoke)
                };
                let ($($a,)*) = self;
                let block_ptr = block as *const _ as *mut _;
                unsafe {
                    invoke(block_ptr $(, $a)*)
                }
            }

            fn invoke_for_concrete_block<R, X: Clone>() ->
                    ConcreteBlockInvoke<($($t,)*), R, X> {
                unsafe extern fn $f<R, X: Clone $(, $t)*>(
                        block_ptr: *mut ConcreteBlock<($($t,)*), R, X>
                        $(, $a: $t)*) -> R {
                    let args = ($($a,)*);
                    let block = &*block_ptr;
                    (block.rust_invoke)(&block.context, args)
                }

                unsafe {
                    mem::transmute($f::<R, X $(, $t)*>)
                }
            }
        }
    );
);

block_args_impl!(concrete_block_invoke_args0);
block_args_impl!(concrete_block_invoke_args1, a: A);
block_args_impl!(concrete_block_invoke_args2, a: A, b: B);
block_args_impl!(concrete_block_invoke_args3, a: A, b: B, c: C);
block_args_impl!(concrete_block_invoke_args4, a: A, b: B, c: C, d: D);
block_args_impl!(concrete_block_invoke_args5, a: A, b: B, c: C, d: D, e: E);
block_args_impl!(concrete_block_invoke_args6, a: A, b: B, c: C, d: D, e: E, f: F);
block_args_impl!(concrete_block_invoke_args7, a: A, b: B, c: C, d: D, e: E, f: F, g: G);
block_args_impl!(concrete_block_invoke_args8, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H);
block_args_impl!(concrete_block_invoke_args9, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I);
block_args_impl!(concrete_block_invoke_args10, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J);
block_args_impl!(concrete_block_invoke_args11, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K);
block_args_impl!(concrete_block_invoke_args12, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K, l: L);

/// An Objective-C block that takes arguments of `A` when called and
/// returns a value of `R`.
#[repr(C)]
pub struct Block<A: BlockArguments, R> {
    isa: *const Class,
    flags: c_int,
    _reserved: c_int,
    invoke: BlockInvoke<A, R>,
}

impl<A: BlockArguments, R> Block<A, R> {
    /// Copy self onto the heap.
    pub fn copy(&self) -> Id<Block<A, R>> {
        unsafe {
            let block = msg_send![self copy] as *mut Block<A, R>;
            Id::from_retained_ptr(block)
        }
    }

    /// Call self with the given arguments.
    pub fn call(&self, args: A) -> R {
        args.call_block(self)
    }
}

impl<A: BlockArguments, R> Message for Block<A, R> { }

/// An Objective-C block whose size is known at compile time and may be
/// constructed on the stack.
#[repr(C)]
pub struct ConcreteBlock<A: BlockArguments, R, C: Clone> {
    base: Block<A, R>,
    descriptor: *const BlockDescriptor<ConcreteBlock<A, R, C>>,
    rust_invoke: fn (&C, A) -> R,
    context: C,
}

impl<A: BlockArguments, R, C: Clone> ConcreteBlock<A, R, C> {
    /// Constructs a `ConcreteBlock` with the given invoke function and context.
    /// When the block is called, it will return the value that results from
    /// calling the invoke function with a reference to its context.
    pub fn new(invoke: fn (&C, A) -> R, context: C) -> ConcreteBlock<A, R, C> {
        let extern_invoke: ConcreteBlockInvoke<A, R, C> =
            BlockArguments::invoke_for_concrete_block();
        ConcreteBlock {
            base: Block {
                isa: &_NSConcreteStackBlock,
                // 1 << 25 = BLOCK_HAS_COPY_DISPOSE
                flags: 1 << 25,
                _reserved: 0,
                invoke: unsafe { mem::transmute(extern_invoke) },
            },
            // TODO: don't leak memory here
            descriptor: unsafe {
                mem::transmute(box BlockDescriptor::<A, R, C>::new())
            },
            rust_invoke: invoke,
            context: context,
        }
    }
}

impl<A: BlockArguments, R, C: Clone> Clone for ConcreteBlock<A, R, C> {
    fn clone(&self) -> ConcreteBlock<A, R, C> {
        ConcreteBlock::new(self.rust_invoke, self.context.clone())
    }
}

impl<A: BlockArguments, R, C: Clone> Deref<Block<A, R>> for ConcreteBlock<A, R, C> {
    fn deref(&self) -> &Block<A, R> {
        &self.base
    }
}

unsafe extern fn block_context_dispose<A: BlockArguments, R, C: Clone>(
        block: &mut ConcreteBlock<A, R, C>) {
    // Read the context from the block and let it drop
    ptr::read(&block.context);
}

unsafe extern fn block_context_copy<A: BlockArguments, R, C: Clone>(
        dst: &mut ConcreteBlock<A, R, C>, src: &ConcreteBlock<A, R, C>) {
    // The src block actually gets memmoved to the destination beforehand,
    // but we'll set the function pointer, too, to be safe.
    ptr::write(&mut dst.rust_invoke, src.rust_invoke);
    ptr::write(&mut dst.context, src.context.clone());
}

#[repr(C)]
struct BlockDescriptor<B> {
    _reserved: c_ulong,
    block_size: c_ulong,
    copy_helper: unsafe extern fn(&mut B, &B),
    dispose_helper: unsafe extern fn(&mut B),
}

impl<A: BlockArguments, R, C: Clone> BlockDescriptor<ConcreteBlock<A, R, C>> {
    fn new() -> BlockDescriptor<ConcreteBlock<A, R, C>> {
        BlockDescriptor {
            _reserved: 0,
            block_size: mem::size_of::<ConcreteBlock<A, R, C>>() as c_ulong,
            copy_helper: block_context_copy::<A, R, C>,
            dispose_helper: block_context_dispose::<A, R, C>,
        }
    }
}

#[cfg(test)]
mod tests {
    use Id;
    use objc_test_utils;
    use super::{Block, ConcreteBlock};

    fn get_int_block() -> &'static Block<(), int> {
        unsafe {
            &*(objc_test_utils::get_int_block() as *const _)
        }
    }

    fn get_int_block_with(i: int) -> Id<Block<(), int>> {
        unsafe {
            let ptr = objc_test_utils::get_int_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn get_add_block() -> &'static Block<(int,), int> {
        unsafe {
            &*(objc_test_utils::get_add_block() as *const _)
        }
    }

    fn get_add_block_with(i: int) -> Id<Block<(int,), int>> {
        unsafe {
            let ptr = objc_test_utils::get_add_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn invoke_int_block(block: &Block<(), int>) -> int {
        let ptr = block as *const _;
        unsafe {
            objc_test_utils::invoke_int_block(ptr as *const _)
        }
    }

    fn invoke_add_block(block: &Block<(int,), int>, a: int) -> int {
        let ptr = block as *const _;
        unsafe {
            objc_test_utils::invoke_add_block(ptr as *const _, a)
        }
    }

    #[test]
    fn test_call_block() {
        let block = get_int_block();
        assert!(block.call(()) == 7);

        let block = get_int_block_with(13);
        assert!(block.call(()) == 13);
    }

    #[test]
    fn test_call_block_args() {
        let block = get_add_block();
        assert!(block.call((2,)) == 9);

        let block = get_add_block_with(13);
        assert!(block.call((2,)) == 15);
    }

    #[test]
    fn test_create_block() {
        fn block_get_int(context: &int, _args: ()) -> int {
            *context
        }

        let block = ConcreteBlock::new(block_get_int, 13);
        let result = invoke_int_block(&*block);
        assert!(result == 13);
    }

    #[test]
    fn test_create_block_args() {
        fn block_add_int(context: &int, (a,): (int,)) -> int {
            a + *context
        }

        let block = ConcreteBlock::new(block_add_int, 5);
        let result = invoke_add_block(&*block, 6);
        assert!(result == 11);
    }

    #[test]
    fn test_concrete_block_clone() {
        fn block_get_string_len(context: &String, _args: ()) -> uint {
            context.len()
        }

        let s = "Hello!".into_string();
        let expected_len = s.len();
        let block = ConcreteBlock::new(block_get_string_len, s);
        assert!(block.call(()) == expected_len);

        let cloned = block.clone();
        assert!(cloned.call(()) == expected_len);

        drop(block);
        assert!(cloned.call(()) == expected_len);
    }

    #[test]
    fn test_concrete_block_copy() {
        fn block_get_string_len(context: &String, _args: ()) -> uint {
            context.len()
        }

        let s = "Hello!".into_string();
        let expected_len = s.len();
        let block = ConcreteBlock::new(block_get_string_len, s);
        assert!(block.call(()) == expected_len);

        let copied = block.copy();
        assert!(copied.call(()) == expected_len);

        drop(block);
        assert!(copied.call(()) == expected_len);
    }
}
<|endoftext|>


<commit_before> #[feature(struct_variant)];
extern mod extra;
use extra::getopts;
use std::rand;
use std::fmt;

fn randrange(min : uint, max : uint) -> uint {
    min + rand::random::<uint>() % (max - min)
}

#[deriving(Eq)]
struct ServerID(uint);

#[deriving(Eq, Ord)]
struct Term(uint);

#[deriving(Ord)]
struct Index(uint);

#[deriving(Ord)]
struct Time(uint);

struct Environment {
    clock: Time,
    network: ~[~Message],
}

impl Environment {
    fn new() -> Environment {
        Environment {
            clock: Time(0),
            network: ~[],
        }
    }
    fn make_time(&self, min: uint, max: uint) -> Time {
        Time(*self.clock + randrange(min, max))
    }
    fn make_network_latency(&self, from: ServerID, to: ServerID) -> Time {
        Time(*self.clock + randrange(2, 5))
    }
    fn multicast(&mut self, from: ServerID, to: &[ServerID], body: &MessageBody) {
        for peer in to.iter() {
            let m = ~Message {
                from: from,
                to: *peer,
                deliver: self.make_network_latency(from, *peer),
                body: *body,
            };
            self.network.push(m);
        }
    }
    fn reply(&mut self, request: &Message, response_body: &MessageBody) {
        let m = ~Message {
            from: request.to,
            to: request.from,
            deliver: self.make_network_latency(request.to, request.from),
            body: *response_body,
        };
        self.network.push(m);
    }
    fn pop_ready_messages(&mut self) -> ~[~Message] {
        let mut all = ~[];
        std::util::swap(&mut all, &mut self.network);
        let (ready, notready) = all.partition(|msg| msg.deliver <= self.clock);
        self.network = notready;
        return ready;
    }
}

struct Message {
    from: ServerID,
    to: ServerID,
    deliver: Time,
    body: MessageBody,
}

enum MessageBody {
    RequestVoteRequest {
        term: Term,
        lastLogTerm: Term,
        lastLogIndex: Index,
    },
    RequestVoteResponse {
        term: Term,
        granted: bool,
    },
    Heartbeat {
        term: Term,
    },
}

impl fmt::Default for Message {
    fn fmt(msg: &Message, f: &mut fmt::Formatter) {
        write!(f.buf, "{} to {}: {}",
               *msg.from,
               *msg.to,
               msg.body);
    }
}

impl fmt::Default for MessageBody {
    fn fmt(b: &MessageBody, f: &mut fmt::Formatter) {
        match *b {
            RequestVoteRequest{term, lastLogTerm, lastLogIndex} =>
                write!(f.buf, "RequestVoteRequest({}, {}, {})",
                       term, lastLogTerm, lastLogIndex),
            RequestVoteResponse{term, granted} =>
                write!(f.buf, "RequestVoteResponse({}, {})",
                       term, granted),
            Heartbeat{term} =>
                write!(f.buf, "Heartbeat({})",
                       term),
        }
    }
}

enum ServerState {
    Follower  { timer: Time },
    Candidate { timer: Time, votes: uint },
    Leader,
}

impl fmt::Default for Time {
    fn fmt(time: &Time, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **time)
    }
}

impl fmt::Default for Term {
    fn fmt(term: &Term, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **term)
    }
}

impl fmt::Default for Index {
    fn fmt(index: &Index, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **index)
    }
}

impl fmt::Default for ServerState {
    fn fmt(state: &ServerState, f: &mut fmt::Formatter) {
        match *state {
            Follower{timer} => write!(f.buf, "Follower({})", timer),
            Candidate{timer, votes} => write!(f.buf, "Candidate({}, {})", timer, votes),
            Leader => write!(f.buf, "Leader"),
        }
    }
}

struct LogEntry {
    term: Term,
    // command: ~str,
}

struct Log(~[LogEntry]);

impl Log {
    fn new() -> Log {
        Log(~[])
    }
}

struct Server {
    id: ServerID,
    peers: ~[ServerID],
    term: Term,
    state: ServerState,
    vote: Option<ServerID>,
    log: Log,
}

impl Server {
    fn new(id : ServerID, peers: ~[ServerID], env: &Environment) -> Server {
        Server {
            id: id,
            peers: peers,
            term: Term(0),
            state: Follower { timer: env.make_time(150, 300) },
            log: Log::new(),
            vote: None,
        }
    }
    /*
    fn reset_timer(&mut self) {
        self.set_timer(Time(*self.env.clock + randrange(150, 300)));
    }
    fn set_timer(&mut self, time: Time) {
        match self.state {
            Follower  {timer : ref mut timer, _} |
            Candidate {timer : ref mut timer, _} =>
                { *timer = time },
            Leader => {},
        }
    }
    */
    fn tick(&mut self, env: &mut Environment) {
        match self.state {
            Follower  {timer, _} |
            Candidate {timer, _} => {
                if timer <= env.clock {
                    self.term = Term(*self.term + 1);
                    self.vote = None;
                    self.state = Candidate {
                        timer: env.make_time(150, 300),
                        votes: 1,
                    };
                    env.multicast(self.id, self.peers, &RequestVoteRequest {
                        term: self.term,
                        lastLogTerm: Term(0),
                        lastLogIndex: Index(0),
                    });
                }
                self.try_become_leader(env);
            },
            Leader => {},
        }
    }
    fn step_down(&mut self, env: &Environment, term: Term) {
        self.term = term;
        self.vote = None;
        self.state = Follower {
            timer: env.make_time(150, 300),
        };
    }

    fn try_become_leader(&mut self, env: &mut Environment) {
        match self.state {
            Follower {_} => {},
            Candidate {votes, _} => {
                if votes > (self.peers.len() + 1) / 2 {
                    self.state = Leader;
                    env.multicast(self.id, self.peers, &Heartbeat {
                        term: self.term,
                    });
                }
            },
            Leader => {},
        }
    }
    fn handle(&mut self, env: &mut Environment, msg: &Message) {
        match msg.body {
            RequestVoteRequest {term, _} => {
                let reply = |granted| {
                    env.reply(msg, &RequestVoteResponse {
                        term: self.term,
                        granted: granted,
                    });
                };
                if term < self.term {
                    reply(false);
                } else {
                    if term > self.term {
                        self.step_down(env, term);
                    }
                    match self.vote {
                        None => {
                            self.vote = Some(msg.from);
                            reply(true);
                        },
                        Some(c) => {
                            reply(c == msg.from);
                        },
                    }
                }
            },
            RequestVoteResponse {term, granted, _} => {
                if term == self.term && granted {
                    match self.state {
                        Follower {_} => {},
                        Candidate {votes: ref mut votes, _} => {
                            *votes += 1;
                        },
                        Leader => {},
                    }
                    self.try_become_leader(env)
                } else if term > self.term {
                    self.step_down(env, term);
                }
            },
            Heartbeat {term} => {
                if term == self.term {
                    self.state = Follower {
                        timer: env.make_time(150, 300),
                    };
                } else if term > self.term {
                    self.step_down(env, term);
                }
            },
        }
    }

    /*
    fn next_event_time(&self) -> Time {
        match self.state {
            Follower  {timer, _} |
            Candidate {timer, _} =>
                timer,
            Leader => Time(!0)
        }
    }
    */
}

impl ToStr for Server {
    fn to_str(&self) -> ~str {
        format!("Server {}: {}", *self.id, self.state)
    }
}

struct Cluster(~[Server]);

impl Cluster {
    fn new(env: &Environment, num_servers: uint) -> Cluster {
        let mut servers = ~[];
        for i in range(0, num_servers) {
            let mut peers = ~[];
            for j in range(0, num_servers) {
                if i != j {
                    peers.push(ServerID(j + 1));
                }
            }
            let s = Server::new(ServerID(i + 1), peers, env);
            servers.push(s);
        }
        return Cluster(servers);
    }

    fn deliver(&mut self, env: &mut Environment, msg: &Message) {
        self[*msg.to - 1].handle(env, msg);
    }
}

fn main() {
    let args = std::os::args();
    let opts = [
        getopts::optflag("h"),
        getopts::optflag("help"),
        getopts::optopt("servers"),
    ];
    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m },
        Err(f) => { fail!(f.to_err_msg()) },
    };
    let usage = || {
        println!("Usage: {} [options]", args[0]);
        println("-h, --help   Print this help message");
        println("--servers=N  Simulate cluster with N servers");
    };
    if matches.opt_present("h") || matches.opt_present("help") {
        usage();
        return;
    }
    if !matches.free.is_empty() {
        usage();
        fail!(format!("Extra arguments: {}", matches.free.len()));
    }
    let num_servers : uint = match matches.opt_str("servers") {
        Some(s) => { match std::from_str::from_str(s) {
            Some(i) => { i },
            None => {
                usage();
                fail!(format!("Couldn't parse number of servers from '{}'", s));
            },
        }},
        None => { 5 },
    };
    let env = &mut Environment::new();
    let mut cluster = Cluster::new(env, num_servers);
    loop {
        env.clock = Time(*env.clock + 1);
        for server in cluster.mut_iter() {
            server.tick(env);
        }
        for msg in env.pop_ready_messages().move_iter() {
            cluster.deliver(env, msg);
        }
        //for server in cluster.iter() {
        //    println(server.to_str());
        //}
        for message in env.network.iter() {
            println!("{}", **message);
        }
        println("");
    }
}

<commit_msg>Add repeated heartbeats<commit_after> #[feature(struct_variant)];
extern mod extra;
use extra::getopts;
use std::rand;
use std::fmt;

fn randrange(min : uint, max : uint) -> uint {
    min + rand::random::<uint>() % (max - min)
}

#[deriving(Eq)]
struct ServerID(uint);

#[deriving(Eq, Ord)]
struct Term(uint);

#[deriving(Ord)]
struct Index(uint);

#[deriving(Ord)]
struct Time(uint);

struct Environment {
    clock: Time,
    network: ~[~Message],
}

impl Environment {
    fn new() -> Environment {
        Environment {
            clock: Time(0),
            network: ~[],
        }
    }
    fn make_time(&self, min: uint, max: uint) -> Time {
        Time(*self.clock + randrange(min, max))
    }
    fn make_network_latency(&self, from: ServerID, to: ServerID) -> Time {
        Time(*self.clock + randrange(2, 5))
    }
    fn multicast(&mut self, from: ServerID, to: &[ServerID], body: &MessageBody) {
        for peer in to.iter() {
            let m = ~Message {
                from: from,
                to: *peer,
                deliver: self.make_network_latency(from, *peer),
                body: *body,
            };
            self.network.push(m);
        }
    }
    fn reply(&mut self, request: &Message, response_body: &MessageBody) {
        let m = ~Message {
            from: request.to,
            to: request.from,
            deliver: self.make_network_latency(request.to, request.from),
            body: *response_body,
        };
        self.network.push(m);
    }
    fn pop_ready_messages(&mut self) -> ~[~Message] {
        let mut all = ~[];
        std::util::swap(&mut all, &mut self.network);
        let (ready, notready) = all.partition(|msg| msg.deliver <= self.clock);
        self.network = notready;
        return ready;
    }
}

struct Message {
    from: ServerID,
    to: ServerID,
    deliver: Time,
    body: MessageBody,
}

enum MessageBody {
    RequestVoteRequest {
        term: Term,
        lastLogTerm: Term,
        lastLogIndex: Index,
    },
    RequestVoteResponse {
        term: Term,
        granted: bool,
    },
    AppendEntriesRequest {
        term: Term,
        seqno: uint,
    },
    AppendEntriesResponse {
        term: Term,
        seqno: uint,
    },
}

impl fmt::Default for Message {
    fn fmt(msg: &Message, f: &mut fmt::Formatter) {
        write!(f.buf, "{} to {}: {}",
               *msg.from,
               *msg.to,
               msg.body);
    }
}

impl fmt::Default for MessageBody {
    fn fmt(b: &MessageBody, f: &mut fmt::Formatter) {
        match *b {
            RequestVoteRequest{term, lastLogTerm, lastLogIndex} =>
                write!(f.buf, "RequestVoteRequest({}, {}, {})",
                       term, lastLogTerm, lastLogIndex),
            RequestVoteResponse{term, granted} =>
                write!(f.buf, "RequestVoteResponse({}, {})",
                       term, granted),
            AppendEntriesRequest{term, seqno} =>
                write!(f.buf, "AppendEntriesRequest({}, {})",
                       term, seqno),
            AppendEntriesResponse{term, seqno} =>
                write!(f.buf, "AppendEntriesResponse({}, {})",
                       term, seqno),
        }
    }
}

enum ServerState {
    Follower  { timer: Time },
    Candidate { timer: Time, votes: uint },
    Leader    { timer: Time, heartbeat_seqno: uint,  acks: uint},
}

impl fmt::Default for Time {
    fn fmt(time: &Time, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **time)
    }
}

impl fmt::Default for Term {
    fn fmt(term: &Term, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **term)
    }
}

impl fmt::Default for Index {
    fn fmt(index: &Index, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **index)
    }
}

impl fmt::Default for ServerState {
    fn fmt(state: &ServerState, f: &mut fmt::Formatter) {
        match *state {
            Follower{timer} => write!(f.buf, "Follower({})", timer),
            Candidate{timer, votes} => write!(f.buf, "Candidate({}, {})", timer, votes),
            Leader{timer, heartbeat_seqno, acks} => write!(f.buf, "Leader({}, {}, {})", timer, heartbeat_seqno, acks),
        }
    }
}

struct LogEntry {
    term: Term,
    // command: ~str,
}

struct Log(~[LogEntry]);

impl Log {
    fn new() -> Log {
        Log(~[])
    }
}

struct Server {
    id: ServerID,
    peers: ~[ServerID],
    term: Term,
    state: ServerState,
    vote: Option<ServerID>,
    log: Log,
}

impl Server {
    fn new(id : ServerID, peers: ~[ServerID], env: &Environment) -> Server {
        Server {
            id: id,
            peers: peers,
            term: Term(0),
            state: Follower { timer: env.make_time(150, 300) },
            log: Log::new(),
            vote: None,
        }
    }
    /*
    fn reset_timer(&mut self) {
        self.set_timer(Time(*self.env.clock + randrange(150, 300)));
    }
    fn set_timer(&mut self, time: Time) {
        match self.state {
            Follower  {timer : ref mut timer, _} |
            Candidate {timer : ref mut timer, _} =>
                { *timer = time },
            Leader => {},
        }
    }
    */
    fn tick(&mut self, env: &mut Environment) {
        let mut need_step_down = false;
        match self.state {
            Follower  {timer, _} |
            Candidate {timer, _} => {
                if timer <= env.clock {
                    self.term = Term(*self.term + 1);
                    self.vote = None;
                    self.state = Candidate {
                        timer: env.make_time(150, 300),
                        votes: 1,
                    };
                    env.multicast(self.id, self.peers, &RequestVoteRequest {
                        term: self.term,
                        lastLogTerm: Term(0),
                        lastLogIndex: Index(0),
                    });
                }
                self.try_become_leader(env);
            },
            Leader {timer, heartbeat_seqno: ref mut heartbeat_seqno, acks: ref mut acks} => {
                if timer <= env.clock {
                    if (*acks > (self.peers.len() + 1) / 2) {
                        *heartbeat_seqno += 1;
                        *acks = 0;
                        env.multicast(self.id, self.peers, &AppendEntriesRequest {
                            term: self.term,
                            seqno: *heartbeat_seqno,
                        });
                    } else {
                        need_step_down = true;
                    }
                }
            },
        }
        if need_step_down {
            self.step_down(env, self.term)
        }
    }
    fn step_down(&mut self, env: &Environment, term: Term) {
        self.term = term;
        self.vote = None;
        self.state = Follower {
            timer: env.make_time(150, 300),
        };
    }

    fn try_become_leader(&mut self, env: &mut Environment) {
        match self.state {
            Follower {_} => {},
            Candidate {votes, _} => {
                if votes > (self.peers.len() + 1) / 2 {
                    self.state = Leader {
                        timer: env.make_time(75, 76),
                        heartbeat_seqno: 0,
                        acks: 1,
                    };
                    env.multicast(self.id, self.peers, &AppendEntriesRequest {
                        term: self.term,
                        seqno: 0,
                    });
                }
            },
            Leader {_} => {},
        }
    }
    fn handle(&mut self, env: &mut Environment, msg: &Message) {
        match msg.body {
            RequestVoteRequest {term, _} => {
                let reply = |granted| {
                    env.reply(msg, &RequestVoteResponse {
                        term: self.term,
                        granted: granted,
                    });
                };
                if term < self.term {
                    reply(false);
                } else {
                    if term > self.term {
                        self.step_down(env, term);
                    }
                    match self.vote {
                        None => {
                            self.vote = Some(msg.from);
                            reply(true);
                        },
                        Some(c) => {
                            reply(c == msg.from);
                        },
                    }
                }
            },
            RequestVoteResponse {term, granted, _} => {
                if term == self.term && granted {
                    match self.state {
                        Follower {_} => {},
                        Candidate {votes: ref mut votes, _} => {
                            *votes += 1;
                        },
                        Leader {_} => {},
                    }
                    self.try_become_leader(env)
                } else if term > self.term {
                    self.step_down(env, term);
                }
            },
            AppendEntriesRequest {term, seqno} => {
                if term == self.term {
                    self.state = Follower {
                        timer: env.make_time(150, 300),
                    };
                } else if term > self.term {
                    self.step_down(env, term);
                }
                env.reply(msg, &AppendEntriesResponse {
                    term: self.term,
                    seqno: seqno,
                });
            },
            AppendEntriesResponse {term, seqno} => {
                if term == self.term {
                    match self.state {
                        Follower {_} | Candidate {_} => {},
                        Leader { heartbeat_seqno, acks: ref mut acks, _ } => {
                            if seqno == heartbeat_seqno {
                                *acks += 1;
                            }
                        },
                    }
                }
            },
        }
    }

    /*
    fn next_event_time(&self) -> Time {
        match self.state {
            Follower  {timer, _} |
            Candidate {timer, _} =>
                timer,
            Leader => Time(!0)
        }
    }
    */
}

impl ToStr for Server {
    fn to_str(&self) -> ~str {
        format!("Server {}: {}", *self.id, self.state)
    }
}

struct Cluster(~[Server]);

impl Cluster {
    fn new(env: &Environment, num_servers: uint) -> Cluster {
        let mut servers = ~[];
        for i in range(0, num_servers) {
            let mut peers = ~[];
            for j in range(0, num_servers) {
                if i != j {
                    peers.push(ServerID(j + 1));
                }
            }
            let s = Server::new(ServerID(i + 1), peers, env);
            servers.push(s);
        }
        return Cluster(servers);
    }

    fn deliver(&mut self, env: &mut Environment, msg: &Message) {
        self[*msg.to - 1].handle(env, msg);
    }
}

fn main() {
    let args = std::os::args();
    let opts = [
        getopts::optflag("h"),
        getopts::optflag("help"),
        getopts::optopt("servers"),
    ];
    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m },
        Err(f) => { fail!(f.to_err_msg()) },
    };
    let usage = || {
        println!("Usage: {} [options]", args[0]);
        println("-h, --help   Print this help message");
        println("--servers=N  Simulate cluster with N servers");
    };
    if matches.opt_present("h") || matches.opt_present("help") {
        usage();
        return;
    }
    if !matches.free.is_empty() {
        usage();
        fail!(format!("Extra arguments: {}", matches.free.len()));
    }
    let num_servers : uint = match matches.opt_str("servers") {
        Some(s) => { match std::from_str::from_str(s) {
            Some(i) => { i },
            None => {
                usage();
                fail!(format!("Couldn't parse number of servers from '{}'", s));
            },
        }},
        None => { 5 },
    };
    let env = &mut Environment::new();
    let mut cluster = Cluster::new(env, num_servers);
    loop {
        env.clock = Time(*env.clock + 1);
        for server in cluster.mut_iter() {
            server.tick(env);
        }
        for msg in env.pop_ready_messages().move_iter() {
            cluster.deliver(env, msg);
        }
        //for server in cluster.iter() {
        //    println(server.to_str());
        //}
        for message in env.network.iter() {
            println!("{}", **message);
        }
    }
}

<|endoftext|>


<commit_before> #[feature(struct_variant)];
extern mod extra;
use extra::getopts;
use std::rand;
use std::fmt;

fn randrange(min : uint, max : uint) -> uint {
    min + rand::random::<uint>() % (max - min)
}

#[deriving(Eq)]
struct ServerID(uint);

#[deriving(Eq, Ord)]
struct Term(uint);

impl fmt::Default for Term {
    fn fmt(term: &Term, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **term)
    }
}

#[deriving(Ord)]
struct Index(uint);

impl fmt::Default for Index {
    fn fmt(index: &Index, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **index)
    }
}

#[deriving(Ord, Clone)]
struct Time(uint);

impl fmt::Default for Time {
    fn fmt(time: &Time, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **time)
    }
}

struct Environment {
    clock: Time,
    network: ~[~Message],
}

impl Environment {
    fn new() -> Environment {
        Environment {
            clock: Time(0),
            network: ~[],
        }
    }
    fn make_time(&self, min: uint, max: uint) -> Time {
        Time(*self.clock + randrange(min, max))
    }
    fn make_network_latency(&self, from: ServerID, to: ServerID) -> Time {
        Time(*self.clock + randrange(2, 5))
    }
    fn multicast(&mut self, from: ServerID, to: &[ServerID], body: &MessageBody) {
        for peer in to.iter() {
            let m = ~Message {
                from: from,
                to: *peer,
                deliver: self.make_network_latency(from, *peer),
                body: *body,
            };
            self.network.push(m);
        }
    }
    fn reply(&mut self, request: &Message, response_body: &MessageBody) {
        let m = ~Message {
            from: request.to,
            to: request.from,
            deliver: self.make_network_latency(request.to, request.from),
            body: *response_body,
        };
        self.network.push(m);
    }
    fn pop_ready_messages(&mut self) -> ~[~Message] {
        if !self.network.iter().any(|msg| msg.deliver <= self.clock) {
            return ~[];
        }
        let mut all = ~[];
        std::util::swap(&mut all, &mut self.network);
        let (ready, notready) = all.partition(|msg| msg.deliver <= self.clock);
        self.network = notready;
        return ready;
    }
}

struct Message {
    from: ServerID,
    to: ServerID,
    deliver: Time,
    body: MessageBody,
}

enum MessageBody {
    RequestVoteRequest {
        term: Term,
        lastLogTerm: Term,
        lastLogIndex: Index,
    },
    RequestVoteResponse {
        term: Term,
        granted: bool,
    },
    AppendEntriesRequest {
        term: Term,
        seqno: uint,
    },
    AppendEntriesResponse {
        term: Term,
        seqno: uint,
    },
}

impl fmt::Default for Message {
    fn fmt(msg: &Message, f: &mut fmt::Formatter) {
        write!(f.buf, "{} to {}: {}",
               *msg.from,
               *msg.to,
               msg.body);
    }
}

impl fmt::Default for MessageBody {
    fn fmt(b: &MessageBody, f: &mut fmt::Formatter) {
        match *b {
            RequestVoteRequest{term, lastLogTerm, lastLogIndex} =>
                write!(f.buf, "RequestVoteRequest(term: {}, lastLogTerm: {}, lastLogIndex: {})",
                       term, lastLogTerm, lastLogIndex),
            RequestVoteResponse{term, granted} =>
                write!(f.buf, "RequestVoteResponse(term: {}, granted: {})",
                       term, granted),
            AppendEntriesRequest{term, seqno} =>
                write!(f.buf, "AppendEntriesRequest(term: {}, seqno: {})",
                       term, seqno),
            AppendEntriesResponse{term, seqno} =>
                write!(f.buf, "AppendEntriesResponse(term: {}, seqno: {})",
                       term, seqno),
        }
    }
}

enum ServerState {
    Follower  { timer: Time },
    Candidate { timer: Time, votes: uint },
    Leader    { timer: Time, heartbeat_seqno: uint, acks: uint, start_time: Time},
}

impl fmt::Default for ServerState {
    fn fmt(state: &ServerState, f: &mut fmt::Formatter) {
        match *state {
            Follower{timer} => {
                write!(f.buf, "Follower(timer: {})",
                       timer)
            },
            Candidate{timer, votes} => {
                write!(f.buf, "Candidate(timer: {}, votes: {})",
                       timer, votes)
            },
            Leader{timer, heartbeat_seqno, acks, start_time} => {
                write!(f.buf, "Leader(timer: {}, seqno: {}, acks: {}, start_time: {})",
                       timer, heartbeat_seqno, acks, start_time)
            },
        }
    }
}

struct LogEntry {
    term: Term,
    // command: ~str,
}

struct Log(~[LogEntry]);

impl Log {
    fn new() -> Log {
        Log(~[])
    }
}

struct Server {
    id: ServerID,
    peers: ~[ServerID],
    term: Term,
    state: ServerState,
    vote: Option<ServerID>,
    log: Log,
}

impl Server {
    fn new(id : ServerID, peers: ~[ServerID], env: &Environment) -> Server {
        Server {
            id: id,
            peers: peers,
            term: Term(0),
            state: Follower { timer: env.make_time(150, 300) },
            log: Log::new(),
            vote: None,
        }
    }
    /*
    fn reset_timer(&mut self) {
        self.set_timer(Time(*self.env.clock + randrange(150, 300)));
    }
    fn set_timer(&mut self, time: Time) {
        match self.state {
            Follower  {timer : ref mut timer, _} |
            Candidate {timer : ref mut timer, _} =>
                { *timer = time },
            Leader => {},
        }
    }
    */
    fn tick(&mut self, env: &mut Environment) {
        let mut need_step_down = false;
        match self.state {
            Follower  {timer, _} |
            Candidate {timer, _} => {
                if timer <= env.clock {
                    self.term = Term(*self.term + 1);
                    self.vote = None;
                    self.state = Candidate {
                        timer: env.make_time(150, 300),
                        votes: 1,
                    };
                    env.multicast(self.id, self.peers, &RequestVoteRequest {
                        term: self.term,
                        lastLogTerm: Term(0),
                        lastLogIndex: Index(0),
                    });
                }
                self.try_become_leader(env);
            },
            Leader {timer: ref mut timer, heartbeat_seqno: ref mut heartbeat_seqno, acks: ref mut acks, _} => {
                if *timer <= env.clock {
                    if (*acks > (self.peers.len() + 1) / 2) {
                        *heartbeat_seqno += 1;
                        *acks = 0;
                        *timer = env.make_time(75, 76);
                        env.multicast(self.id, self.peers, &AppendEntriesRequest {
                            term: self.term,
                            seqno: *heartbeat_seqno,
                        });
                    } else {
                        need_step_down = true;
                    }
                }
            },
        }
        if need_step_down {
            println!("Failed to maintain leadership: {}", *self);
            self.step_down(env, self.term)
        }
    }
    fn step_down(&mut self, env: &Environment, term: Term) {
        self.term = term;
        self.vote = None;
        self.state = Follower {
            timer: env.make_time(150, 300),
        };
    }

    fn try_become_leader(&mut self, env: &mut Environment) {
        match self.state {
            Follower {_} => {},
            Candidate {votes, _} => {
                if votes > (self.peers.len() + 1) / 2 {
                    self.state = Leader {
                        timer: env.make_time(75, 76),
                        heartbeat_seqno: 0,
                        acks: 1,
                        start_time: env.clock,
                    };
                    env.multicast(self.id, self.peers, &AppendEntriesRequest {
                        term: self.term,
                        seqno: 0,
                    });
                }
            },
            Leader {_} => {},
        }
    }

    fn stable_leader_start_time(&self) -> Option<Time> {
        match self.state {
            Leader {heartbeat_seqno, start_time, _} if heartbeat_seqno > 4 => Some(start_time),
            _ => None,
        }
    }

    fn handle(&mut self, env: &mut Environment, msg: &Message) {
        match msg.body {
            RequestVoteRequest {term, _} => {
                let reply = |granted| {
                    env.reply(msg, &RequestVoteResponse {
                        term: self.term,
                        granted: granted,
                    });
                };
                if term < self.term {
                    reply(false);
                } else {
                    if term > self.term {
                        self.step_down(env, term);
                    }
                    match self.vote {
                        None => {
                            self.vote = Some(msg.from);
                            reply(true);
                        },
                        Some(c) => {
                            reply(c == msg.from);
                        },
                    }
                }
            },
            RequestVoteResponse {term, granted, _} => {
                if term == self.term && granted {
                    match self.state {
                        Follower {_} => {},
                        Candidate {votes: ref mut votes, _} => {
                            *votes += 1;
                        },
                        Leader {_} => {},
                    }
                    self.try_become_leader(env)
                } else if term > self.term {
                    self.step_down(env, term);
                }
            },
            AppendEntriesRequest {term, seqno} => {
                if term == self.term {
                    self.state = Follower {
                        timer: env.make_time(150, 300),
                    };
                } else if term > self.term {
                    self.step_down(env, term);
                }
                env.reply(msg, &AppendEntriesResponse {
                    term: self.term,
                    seqno: seqno,
                });
            },
            AppendEntriesResponse {term, seqno} => {
                if term == self.term {
                    match self.state {
                        Follower {_} | Candidate {_} => {},
                        Leader { heartbeat_seqno, acks: ref mut acks, _ } => {
                            if seqno == heartbeat_seqno {
                                *acks += 1;
                            }
                        },
                    }
                }
            },
        }
    }

    /*
    fn next_event_time(&self) -> Time {
        match self.state {
            Follower  {timer, _} |
            Candidate {timer, _} =>
                timer,
            Leader => Time(!0)
        }
    }
    */
}

impl fmt::Default for Server {
    fn fmt(server: &Server, f: &mut fmt::Formatter) {
        write!(f.buf, "Server {}: {}", *server.id, server.state)
    }
}

struct Cluster(~[Server]);

impl Cluster {
    fn new(env: &Environment, num_servers: uint) -> Cluster {
        let mut servers = ~[];
        for i in range(0, num_servers) {
            let mut peers = ~[];
            for j in range(0, num_servers) {
                if i != j {
                    peers.push(ServerID(j + 1));
                }
            }
            let s = Server::new(ServerID(i + 1), peers, env);
            servers.push(s);
        }
        return Cluster(servers);
    }

    fn deliver(&mut self, env: &mut Environment, msg: &Message) {
        self[*msg.to - 1].handle(env, msg);
    }
}

fn main() {
    let args = std::os::args();
    let opts = [
        getopts::optflag("h"),
        getopts::optflag("help"),
        getopts::optopt("samples"),
        getopts::optopt("servers"),
        getopts::optopt("tasks"),
    ];
    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m },
        Err(f) => { fail!(f.to_err_msg()) },
    };
    let usage = || {
        println!("Usage: {} [options]", args[0]);
        println("-h, --help    Print this help message");
        println("--samples=N   Number of simulations (default 100,000)");
        println("--servers=N   Simulate cluster with N servers");
        println("--tasks=N     Number of parallel jobs (default 1)");
    };
    if matches.opt_present("h") || matches.opt_present("help") {
        usage();
        return;
    }
    if !matches.free.is_empty() {
        usage();
        fail!(format!("Extra arguments: {}", matches.free.len()));
    }
    let num_samples : uint = match matches.opt_str("samples") {
        Some(s) => { match std::from_str::from_str(s) {
            Some(i) => { i },
            None => {
                usage();
                fail!(format!("Couldn't parse number of samples from '{}'", s));
            },
        }},
        None => { 100000 },
    };
    let num_servers : uint = match matches.opt_str("servers") {
        Some(s) => { match std::from_str::from_str(s) {
            Some(i) => { i },
            None => {
                usage();
                fail!(format!("Couldn't parse number of servers from '{}'", s));
            },
        }},
        None => { 5 },
    };
    let num_tasks : uint = match matches.opt_str("tasks") {
        Some(s) => { match std::from_str::from_str(s) {
            Some(i) => { i },
            None => {
                usage();
                fail!(format!("Couldn't parse number of tasks from '{}'", s));
            },
        }},
        None => { 1 },
    };

    let (port, chan): (Port<~[Time]>, Chan<~[Time]>) = stream();
    let chan = std::comm::SharedChan::new(chan);
    for tid in range(0, num_tasks) {
        let child_chan = chan.clone();
        do spawn || {
            let mut samples = ~[];
            let n = if tid == 0 {
                // get the odd one left over
                num_samples - (num_samples / num_tasks) * (num_tasks - 1)
            } else {
                num_samples / num_tasks
            };
            samples.reserve(n);
            do n.times {
                samples.push(simulate(num_servers))
            }
            extra::sort::tim_sort(samples);
            child_chan.send(samples);
        }
    }

    let mut samples = ~[];
    samples.reserve(num_samples);
    do num_tasks.times {
        samples.push_all(port.recv());
    }
    extra::sort::tim_sort(samples);

    println!("Across {} trials:", samples.len());
    println!("Min:    {} ticks", samples[0]);
    println!("Median: {} ticks", samples[samples.len()/2]);
    println!("Max:    {} ticks", samples[samples.len()-1]);
}

fn simulate(num_servers: uint) -> Time {
    let env = &mut Environment::new();
    let mut cluster = Cluster::new(env, num_servers);
    let mut end = None;
    while end.is_none() {
        env.clock = Time(*env.clock + 1);
        for server in cluster.mut_iter() {
            server.tick(env);
        }
        for msg in env.pop_ready_messages().move_iter() {
            cluster.deliver(env, msg);
        }
        for server in cluster.iter() {
            end = server.stable_leader_start_time();
            if end.is_some() {
                break;
            }
        }
        /*
        println!("Tick: {}", env.clock);
        for server in cluster.iter() {
            println!("{}", *server);
        }
        for message in env.network.iter() {
            println!("{}", **message);
        }
        println!("");
        */
    }
    return end.unwrap();
}

<commit_msg>Start TimingPolicy<commit_after> #[feature(struct_variant)];
extern mod extra;
use extra::getopts;
use std::rand;
use std::fmt;

fn randrange(min : uint, max : uint) -> uint {
    min + rand::random::<uint>() % (max - min)
}

#[deriving(Eq)]
struct ServerID(uint);

#[deriving(Eq, Ord)]
struct Term(uint);

impl fmt::Default for Term {
    fn fmt(term: &Term, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **term)
    }
}

#[deriving(Ord)]
struct Index(uint);

impl fmt::Default for Index {
    fn fmt(index: &Index, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **index)
    }
}

#[deriving(Ord, Clone)]
struct Time(uint);

impl Add<uint, Time> for Time {
    fn add(&self, rhs: &uint) -> Time {
      Time(**self + *rhs)
  }
}

impl fmt::Default for Time {
    fn fmt(time: &Time, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **time)
    }
}

struct Environment {
    clock: Time,
    network: ~[~Message],
    timing: ~TimingPolicy,
}

trait TimingPolicy {
    fn network_latency(&self, from: ServerID, to: ServerID) -> uint;
}

struct UniformTimingPolicy {
    latency_range : (uint, uint),
}

impl TimingPolicy for UniformTimingPolicy {
    fn network_latency(&self, from: ServerID, to: ServerID) -> uint {
        match self.latency_range {
          (min, max) => randrange(min, max)
        }
    }
}

impl Environment {
    fn new(timing: ~TimingPolicy) -> Environment {
        Environment {
            clock: Time(0),
            network: ~[],
            timing: timing,
        }
    }
    fn make_time(&self, min: uint, max: uint) -> Time {
        self.clock + randrange(min, max)
    }
    fn multicast(&mut self, from: ServerID, to: &[ServerID], body: &MessageBody) {
        for peer in to.iter() {
            let m = ~Message {
                from: from,
                to: *peer,
                deliver: self.clock + self.timing.network_latency(from, *peer),
                body: *body,
            };
            self.network.push(m);
        }
    }
    fn reply(&mut self, request: &Message, response_body: &MessageBody) {
        let m = ~Message {
            from: request.to,
            to: request.from,
            deliver: self.clock + self.timing.network_latency(request.to, request.from),
            body: *response_body,
        };
        self.network.push(m);
    }
    fn pop_ready_messages(&mut self) -> ~[~Message] {
        if !self.network.iter().any(|msg| msg.deliver <= self.clock) {
            return ~[];
        }
        let mut all = ~[];
        std::util::swap(&mut all, &mut self.network);
        let (ready, notready) = all.partition(|msg| msg.deliver <= self.clock);
        self.network = notready;
        return ready;
    }
}

struct Message {
    from: ServerID,
    to: ServerID,
    deliver: Time,
    body: MessageBody,
}

enum MessageBody {
    RequestVoteRequest {
        term: Term,
        lastLogTerm: Term,
        lastLogIndex: Index,
    },
    RequestVoteResponse {
        term: Term,
        granted: bool,
    },
    AppendEntriesRequest {
        term: Term,
        seqno: uint,
    },
    AppendEntriesResponse {
        term: Term,
        seqno: uint,
    },
}

impl fmt::Default for Message {
    fn fmt(msg: &Message, f: &mut fmt::Formatter) {
        write!(f.buf, "{} to {}: {}",
               *msg.from,
               *msg.to,
               msg.body);
    }
}

impl fmt::Default for MessageBody {
    fn fmt(b: &MessageBody, f: &mut fmt::Formatter) {
        match *b {
            RequestVoteRequest{term, lastLogTerm, lastLogIndex} =>
                write!(f.buf, "RequestVoteRequest(term: {}, lastLogTerm: {}, lastLogIndex: {})",
                       term, lastLogTerm, lastLogIndex),
            RequestVoteResponse{term, granted} =>
                write!(f.buf, "RequestVoteResponse(term: {}, granted: {})",
                       term, granted),
            AppendEntriesRequest{term, seqno} =>
                write!(f.buf, "AppendEntriesRequest(term: {}, seqno: {})",
                       term, seqno),
            AppendEntriesResponse{term, seqno} =>
                write!(f.buf, "AppendEntriesResponse(term: {}, seqno: {})",
                       term, seqno),
        }
    }
}

enum ServerState {
    Follower  { timer: Time },
    Candidate { timer: Time, votes: uint },
    Leader    { timer: Time, heartbeat_seqno: uint, acks: uint, start_time: Time},
}

impl fmt::Default for ServerState {
    fn fmt(state: &ServerState, f: &mut fmt::Formatter) {
        match *state {
            Follower{timer} => {
                write!(f.buf, "Follower(timer: {})",
                       timer)
            },
            Candidate{timer, votes} => {
                write!(f.buf, "Candidate(timer: {}, votes: {})",
                       timer, votes)
            },
            Leader{timer, heartbeat_seqno, acks, start_time} => {
                write!(f.buf, "Leader(timer: {}, seqno: {}, acks: {}, start_time: {})",
                       timer, heartbeat_seqno, acks, start_time)
            },
        }
    }
}

struct LogEntry {
    term: Term,
    // command: ~str,
}

struct Log(~[LogEntry]);

impl Log {
    fn new() -> Log {
        Log(~[])
    }
}

struct Server {
    id: ServerID,
    peers: ~[ServerID],
    term: Term,
    state: ServerState,
    vote: Option<ServerID>,
    log: Log,
}

impl Server {
    fn new(id : ServerID, peers: ~[ServerID], env: &Environment) -> Server {
        Server {
            id: id,
            peers: peers,
            term: Term(0),
            state: Follower { timer: env.make_time(150, 300) },
            log: Log::new(),
            vote: None,
        }
    }
    /*
    fn reset_timer(&mut self) {
        self.set_timer(Time(*self.env.clock + randrange(150, 300)));
    }
    fn set_timer(&mut self, time: Time) {
        match self.state {
            Follower  {timer : ref mut timer, _} |
            Candidate {timer : ref mut timer, _} =>
                { *timer = time },
            Leader => {},
        }
    }
    */
    fn tick(&mut self, env: &mut Environment) {
        let mut need_step_down = false;
        match self.state {
            Follower  {timer, _} |
            Candidate {timer, _} => {
                if timer <= env.clock {
                    self.term = Term(*self.term + 1);
                    self.vote = None;
                    self.state = Candidate {
                        timer: env.make_time(150, 300),
                        votes: 1,
                    };
                    env.multicast(self.id, self.peers, &RequestVoteRequest {
                        term: self.term,
                        lastLogTerm: Term(0),
                        lastLogIndex: Index(0),
                    });
                }
                self.try_become_leader(env);
            },
            Leader {timer: ref mut timer, heartbeat_seqno: ref mut heartbeat_seqno, acks: ref mut acks, _} => {
                if *timer <= env.clock {
                    if (*acks > (self.peers.len() + 1) / 2) {
                        *heartbeat_seqno += 1;
                        *acks = 0;
                        *timer = env.make_time(75, 76);
                        env.multicast(self.id, self.peers, &AppendEntriesRequest {
                            term: self.term,
                            seqno: *heartbeat_seqno,
                        });
                    } else {
                        need_step_down = true;
                    }
                }
            },
        }
        if need_step_down {
            println!("Failed to maintain leadership: {}", *self);
            self.step_down(env, self.term)
        }
    }
    fn step_down(&mut self, env: &Environment, term: Term) {
        self.term = term;
        self.vote = None;
        self.state = Follower {
            timer: env.make_time(150, 300),
        };
    }

    fn try_become_leader(&mut self, env: &mut Environment) {
        match self.state {
            Follower {_} => {},
            Candidate {votes, _} => {
                if votes > (self.peers.len() + 1) / 2 {
                    self.state = Leader {
                        timer: env.make_time(75, 76),
                        heartbeat_seqno: 0,
                        acks: 1,
                        start_time: env.clock,
                    };
                    env.multicast(self.id, self.peers, &AppendEntriesRequest {
                        term: self.term,
                        seqno: 0,
                    });
                }
            },
            Leader {_} => {},
        }
    }

    fn stable_leader_start_time(&self) -> Option<Time> {
        match self.state {
            Leader {heartbeat_seqno, start_time, _} if heartbeat_seqno > 4 => Some(start_time),
            _ => None,
        }
    }

    fn handle(&mut self, env: &mut Environment, msg: &Message) {
        match msg.body {
            RequestVoteRequest {term, _} => {
                let reply = |granted| {
                    env.reply(msg, &RequestVoteResponse {
                        term: self.term,
                        granted: granted,
                    });
                };
                if term < self.term {
                    reply(false);
                } else {
                    if term > self.term {
                        self.step_down(env, term);
                    }
                    match self.vote {
                        None => {
                            self.vote = Some(msg.from);
                            reply(true);
                        },
                        Some(c) => {
                            reply(c == msg.from);
                        },
                    }
                }
            },
            RequestVoteResponse {term, granted, _} => {
                if term == self.term && granted {
                    match self.state {
                        Follower {_} => {},
                        Candidate {votes: ref mut votes, _} => {
                            *votes += 1;
                        },
                        Leader {_} => {},
                    }
                    self.try_become_leader(env)
                } else if term > self.term {
                    self.step_down(env, term);
                }
            },
            AppendEntriesRequest {term, seqno} => {
                if term == self.term {
                    self.state = Follower {
                        timer: env.make_time(150, 300),
                    };
                } else if term > self.term {
                    self.step_down(env, term);
                }
                env.reply(msg, &AppendEntriesResponse {
                    term: self.term,
                    seqno: seqno,
                });
            },
            AppendEntriesResponse {term, seqno} => {
                if term == self.term {
                    match self.state {
                        Follower {_} | Candidate {_} => {},
                        Leader { heartbeat_seqno, acks: ref mut acks, _ } => {
                            if seqno == heartbeat_seqno {
                                *acks += 1;
                            }
                        },
                    }
                }
            },
        }
    }

    /*
    fn next_event_time(&self) -> Time {
        match self.state {
            Follower  {timer, _} |
            Candidate {timer, _} =>
                timer,
            Leader => Time(!0)
        }
    }
    */
}

impl fmt::Default for Server {
    fn fmt(server: &Server, f: &mut fmt::Formatter) {
        write!(f.buf, "Server {}: {}", *server.id, server.state)
    }
}

struct Cluster(~[Server]);

impl Cluster {
    fn new(env: &Environment, num_servers: uint) -> Cluster {
        let mut servers = ~[];
        for i in range(0, num_servers) {
            let mut peers = ~[];
            for j in range(0, num_servers) {
                if i != j {
                    peers.push(ServerID(j + 1));
                }
            }
            let s = Server::new(ServerID(i + 1), peers, env);
            servers.push(s);
        }
        return Cluster(servers);
    }

    fn deliver(&mut self, env: &mut Environment, msg: &Message) {
        self[*msg.to - 1].handle(env, msg);
    }
}

fn main() {
    let args = std::os::args();
    let opts = [
        getopts::optflag("h"),
        getopts::optflag("help"),
        getopts::optopt("samples"),
        getopts::optopt("servers"),
        getopts::optopt("tasks"),
    ];
    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m },
        Err(f) => { fail!(f.to_err_msg()) },
    };
    let usage = || {
        println!("Usage: {} [options]", args[0]);
        println("-h, --help    Print this help message");
        println("--samples=N   Number of simulations (default 100,000)");
        println("--servers=N   Simulate cluster with N servers");
        println("--tasks=N     Number of parallel jobs (default 1)");
    };
    if matches.opt_present("h") || matches.opt_present("help") {
        usage();
        return;
    }
    if !matches.free.is_empty() {
        usage();
        fail!(format!("Extra arguments: {}", matches.free.len()));
    }
    let num_samples : uint = match matches.opt_str("samples") {
        Some(s) => { match std::from_str::from_str(s) {
            Some(i) => { i },
            None => {
                usage();
                fail!(format!("Couldn't parse number of samples from '{}'", s));
            },
        }},
        None => { 100000 },
    };
    let num_servers : uint = match matches.opt_str("servers") {
        Some(s) => { match std::from_str::from_str(s) {
            Some(i) => { i },
            None => {
                usage();
                fail!(format!("Couldn't parse number of servers from '{}'", s));
            },
        }},
        None => { 5 },
    };
    let num_tasks : uint = match matches.opt_str("tasks") {
        Some(s) => { match std::from_str::from_str(s) {
            Some(i) => { i },
            None => {
                usage();
                fail!(format!("Couldn't parse number of tasks from '{}'", s));
            },
        }},
        None => { 1 },
    };

    let (port, chan): (Port<~[Time]>, Chan<~[Time]>) = stream();
    let chan = std::comm::SharedChan::new(chan);
    for tid in range(0, num_tasks) {
        let child_chan = chan.clone();
        do spawn || {
            let mut samples = ~[];
            let n = if tid == 0 {
                // get the odd one left over
                num_samples - (num_samples / num_tasks) * (num_tasks - 1)
            } else {
                num_samples / num_tasks
            };
            samples.reserve(n);
            do n.times {
                samples.push(simulate(num_servers))
            }
            extra::sort::tim_sort(samples);
            child_chan.send(samples);
        }
    }

    let mut samples = ~[];
    samples.reserve(num_samples);
    do num_tasks.times {
        samples.push_all(port.recv());
    }
    extra::sort::tim_sort(samples);

    println!("Across {} trials:", samples.len());
    println!("Min:    {} ticks", samples[0]);
    println!("Median: {} ticks", samples[samples.len()/2]);
    println!("Max:    {} ticks", samples[samples.len()-1]);
}

fn simulate(num_servers: uint) -> Time {
    let env = &mut Environment::new(~UniformTimingPolicy {
                    latency_range: (2, 5),
                } as ~TimingPolicy);
    let mut cluster = Cluster::new(env, num_servers);
    let mut end = None;
    while end.is_none() {
        env.clock = Time(*env.clock + 1);
        for server in cluster.mut_iter() {
            server.tick(env);
        }
        for msg in env.pop_ready_messages().move_iter() {
            cluster.deliver(env, msg);
        }
        for server in cluster.iter() {
            end = server.stable_leader_start_time();
            if end.is_some() {
                break;
            }
        }
        /*
        println!("Tick: {}", env.clock);
        for server in cluster.iter() {
            println!("{}", *server);
        }
        for message in env.network.iter() {
            println!("{}", **message);
        }
        println!("");
        */
    }
    return end.unwrap();
}

<|endoftext|>