<commit_before>use std::{error, fmt};

use gl;
use gl::{AttribLoc, Enum, UnifLoc};

pub struct VertexArray<'a> {
  pub data: &'a [f32],
  pub components: u32,
  pub stride: u32,
}

pub struct Program {
  id: gl::Program,
  vertex_shader: Shader,
  fragment_shader: Shader,
  pub mvp_matrix: UnifLoc,
  #[allow(dead_code)]
  position: AttribLoc,
  pub texture_unit: UnifLoc,
  #[allow(dead_code)]
  texture_coord: AttribLoc,
}

impl Drop for Program {
  fn drop(&mut self) {
    gl::disable_vertex_attrib_array(self.position);
    gl::disable_vertex_attrib_array(self.texture_coord);
    gl::detach_shader(self.id, self.vertex_shader.id);
    println!("***** Vertex shader(id: {}) detached", self.vertex_shader.id);
    gl::detach_shader(self.id, self.fragment_shader.id);
    println!("***** Fragment shader(id: {}) detached", self.fragment_shader.id);
    gl::delete_program(self.id);
    println!("***** Program(id: {}) deleted", self.id);
  }
}

impl Program {
  pub fn new() -> Result<Program, GlError> {
    match Shader::new(VERTEX_SHADER, gl::VERTEX_SHADER) {
      Err(e) => Err(e),
      Ok(vs) => {
        println!("***** Vertex shader(id: {}) compiled", vs.id);
        match Shader::new(FRAGMENT_SHADER, gl::FRAGMENT_SHADER) {
          Err(e) => Err(e),
          Ok(fs) => {
            println!("***** Fragment shader(id: {}) compiled", vs.id);
            Program::new_from_shaders(vs, fs)
          }
        }
      }
    }
  }

  fn new_from_shaders(vertex_shader: Shader, fragment_shader: Shader) ->
    Result<Program, GlError> {

    let id = match gl::create_program() {
      Ok(id) => id,
      Err(e) => return Err(GlError::Generic(format!("gl::create_program() failed: {:?}", e))),
    };

    gl::attach_shader(id, vertex_shader.id);
    gl::attach_shader(id, fragment_shader.id);
    gl::bind_attrib_location(id, 0, "a_Position");
    gl::bind_attrib_location(id, 1, "a_TextureCoord");
    gl::link_program(id);

    if !gl::get_link_status(id) {
      let info_log = gl::get_program_info_log(id);
      return Err(GlError::Generic(format!("Linking program failed: {}", info_log)));
    }

    let mvp_matrix = match gl::get_uniform_location(id, "u_MVPMatrix") {
      Ok(l) => l,
      Err(e) => return Err(GlError::Generic(format!("gl::get_uniform_location(\"u_MVPMatrix\") failed: {:?}", e))),
    };
    let position = match gl::get_attrib_location(id, "a_Position") {
      Ok(l) => l,
      Err(e) => return Err(GlError::Generic(format!("gl::get_attrib_location(\"a_Position\") failed: {:?}", e))),
    };
    let texture_unit = match gl::get_uniform_location(id, "u_TextureUnit") {
      Ok(l) => l,
      Err(e) => return Err(GlError::Generic(format!("gl::get_uniform_location(\"u_TextureUnit\") failed: {:?}", e))),
    };
    let texture_coord = match gl::get_attrib_location(id, "a_TextureCoord") {
      Ok(l) => l,
      Err(e) => return Err(GlError::Generic(format!("gl::get_attrib_location(\"a_TextureCoord\") failed: {:?}", e))),
    };
    gl::use_program(id);

    println!("***** Using program(id: {})", id);

    let program = Program {
      id: id,
      vertex_shader: vertex_shader,
      fragment_shader: fragment_shader,
      mvp_matrix: mvp_matrix,
      position: position,
      texture_unit: texture_unit,
      texture_coord: texture_coord,
    };
    Ok(program)
  }

  /// Set the vertex attributes for position and texture coordinate.
  pub fn set_vertices(&self, vertex_coords: &VertexArray, texture_coords: &VertexArray) {
    gl::vertex_attrib_pointer_f32(self.position, vertex_coords.components as i32,
      vertex_coords.stride as i32, vertex_coords.data);
    gl::enable_vertex_attrib_array(self.position);

    gl::vertex_attrib_pointer_f32(self.texture_coord, texture_coords.components as i32,
      texture_coords.stride as i32, texture_coords.data);
    gl::enable_vertex_attrib_array(self.texture_coord);
  }
}

pub struct Shader {
  id: gl::Shader,
}

impl Drop for Shader {
  fn drop(&mut self) {
    gl::delete_shader(self.id);
    println!("***** Shader(id: {}) deleted", self.id);
  }
}

impl Shader {
  fn new(shader_string: &str, shader_type: Enum) -> Result<Shader, GlError> {
    let id = {
      match gl::create_shader(shader_type) {
        Ok(id) => id,
        Err(_) => return Err(GlError::Generic("gl::create_shader() failed".to_string())),
      }
    };
    let shader = Shader { id: id };
    gl::shader_source(shader.id, shader_string);
    gl::compile_shader(shader.id);

    if !gl::get_compile_status(shader.id) {
      let info_log = gl::get_shader_info_log(shader.id);
      return Err(GlError::Generic(format!("Compiling shader {} failed: {}", shader_type, info_log)));
    }
    Ok(shader)
  }
}

/// Error that can happen while calling GL.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum GlError {
  Generic(String),
  ShaderCompileFailed(String),
  ProgramLinkFailed(String),
  NotSupported,
}

impl GlError {
  fn to_string(&self) -> &str {
    match *self {
      GlError::Generic(ref text) => &text,
      GlError::ShaderCompileFailed(ref text) => &text,
      GlError::ProgramLinkFailed(ref text) => &text,
      GlError::NotSupported => "Some of the requested attributes are not supported",
    }
  }
}

impl fmt::Display for GlError {
  fn fmt(&self, formatter: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    formatter.write_str(self.to_string())
  }
}

impl error::Error for GlError {
  fn description(&self) -> &str {
    self.to_string()
  }
}

#[cfg(target_os = "android")]
static VERTEX_SHADER: &'static str = include_str!("vertex_shader.gles.glsl");
#[cfg(target_os = "android")]
static FRAGMENT_SHADER: &'static str = include_str!("fragment_shader.gles.glsl");

#[cfg(target_os = "linux")]
static VERTEX_SHADER: &'static str = include_str!("vertex_shader.mesa.glsl");
#[cfg(target_os = "linux")]
static FRAGMENT_SHADER: &'static str = include_str!("fragment_shader.mesa.glsl");
<commit_msg>Delete debugging printouts on GL handle creation and deletion.<commit_after>use std::{error, fmt};

use gl;
use gl::{AttribLoc, Enum, UnifLoc};

pub struct VertexArray<'a> {
  pub data: &'a [f32],
  pub components: u32,
  pub stride: u32,
}

pub struct Program {
  id: gl::Program,
  vertex_shader: Shader,
  fragment_shader: Shader,
  pub mvp_matrix: UnifLoc,
  #[allow(dead_code)]
  position: AttribLoc,
  pub texture_unit: UnifLoc,
  #[allow(dead_code)]
  texture_coord: AttribLoc,
}

impl Drop for Program {
  fn drop(&mut self) {
    gl::disable_vertex_attrib_array(self.position);
    gl::disable_vertex_attrib_array(self.texture_coord);
    gl::detach_shader(self.id, self.vertex_shader.id);
    gl::detach_shader(self.id, self.fragment_shader.id);
    gl::delete_program(self.id);
  }
}

impl Program {
  pub fn new() -> Result<Program, GlError> {
    match Shader::new(VERTEX_SHADER, gl::VERTEX_SHADER) {
      Err(e) => Err(e),
      Ok(vs) => {
        match Shader::new(FRAGMENT_SHADER, gl::FRAGMENT_SHADER) {
          Err(e) => Err(e),
          Ok(fs) => Program::new_from_shaders(vs, fs),
        }
      }
    }
  }

  fn new_from_shaders(vertex_shader: Shader, fragment_shader: Shader) ->
    Result<Program, GlError> {

    let id = match gl::create_program() {
      Ok(id) => id,
      Err(e) => return Err(GlError::Generic(format!("gl::create_program() failed: {:?}", e))),
    };

    gl::attach_shader(id, vertex_shader.id);
    gl::attach_shader(id, fragment_shader.id);
    gl::bind_attrib_location(id, 0, "a_Position");
    gl::bind_attrib_location(id, 1, "a_TextureCoord");
    gl::link_program(id);

    if !gl::get_link_status(id) {
      let info_log = gl::get_program_info_log(id);
      return Err(GlError::Generic(format!("Linking program failed: {}", info_log)));
    }

    let mvp_matrix = match gl::get_uniform_location(id, "u_MVPMatrix") {
      Ok(l) => l,
      Err(e) => return Err(GlError::Generic(format!("gl::get_uniform_location(\"u_MVPMatrix\") failed: {:?}", e))),
    };
    let position = match gl::get_attrib_location(id, "a_Position") {
      Ok(l) => l,
      Err(e) => return Err(GlError::Generic(format!("gl::get_attrib_location(\"a_Position\") failed: {:?}", e))),
    };
    let texture_unit = match gl::get_uniform_location(id, "u_TextureUnit") {
      Ok(l) => l,
      Err(e) => return Err(GlError::Generic(format!("gl::get_uniform_location(\"u_TextureUnit\") failed: {:?}", e))),
    };
    let texture_coord = match gl::get_attrib_location(id, "a_TextureCoord") {
      Ok(l) => l,
      Err(e) => return Err(GlError::Generic(format!("gl::get_attrib_location(\"a_TextureCoord\") failed: {:?}", e))),
    };
    gl::use_program(id);

    let program = Program {
      id: id,
      vertex_shader: vertex_shader,
      fragment_shader: fragment_shader,
      mvp_matrix: mvp_matrix,
      position: position,
      texture_unit: texture_unit,
      texture_coord: texture_coord,
    };
    Ok(program)
  }

  /// Set the vertex attributes for position and texture coordinate.
  pub fn set_vertices(&self, vertex_coords: &VertexArray, texture_coords: &VertexArray) {
    gl::vertex_attrib_pointer_f32(self.position, vertex_coords.components as i32,
      vertex_coords.stride as i32, vertex_coords.data);
    gl::enable_vertex_attrib_array(self.position);

    gl::vertex_attrib_pointer_f32(self.texture_coord, texture_coords.components as i32,
      texture_coords.stride as i32, texture_coords.data);
    gl::enable_vertex_attrib_array(self.texture_coord);
  }
}

pub struct Shader {
  id: gl::Shader,
}

impl Drop for Shader {
  fn drop(&mut self) {
    gl::delete_shader(self.id);
  }
}

impl Shader {
  fn new(shader_string: &str, shader_type: Enum) -> Result<Shader, GlError> {
    let id = {
      match gl::create_shader(shader_type) {
        Ok(id) => id,
        Err(_) => return Err(GlError::Generic("gl::create_shader() failed".to_string())),
      }
    };
    let shader = Shader { id: id };
    gl::shader_source(shader.id, shader_string);
    gl::compile_shader(shader.id);

    if !gl::get_compile_status(shader.id) {
      let info_log = gl::get_shader_info_log(shader.id);
      return Err(GlError::Generic(format!("Compiling shader {} failed: {}", shader_type, info_log)));
    }
    Ok(shader)
  }
}

/// Error that can happen while calling GL.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum GlError {
  Generic(String),
  ShaderCompileFailed(String),
  ProgramLinkFailed(String),
  NotSupported,
}

impl GlError {
  fn to_string(&self) -> &str {
    match *self {
      GlError::Generic(ref text) => &text,
      GlError::ShaderCompileFailed(ref text) => &text,
      GlError::ProgramLinkFailed(ref text) => &text,
      GlError::NotSupported => "Some of the requested attributes are not supported",
    }
  }
}

impl fmt::Display for GlError {
  fn fmt(&self, formatter: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    formatter.write_str(self.to_string())
  }
}

impl error::Error for GlError {
  fn description(&self) -> &str {
    self.to_string()
  }
}

#[cfg(target_os = "android")]
static VERTEX_SHADER: &'static str = include_str!("vertex_shader.gles.glsl");
#[cfg(target_os = "android")]
static FRAGMENT_SHADER: &'static str = include_str!("fragment_shader.gles.glsl");

#[cfg(target_os = "linux")]
static VERTEX_SHADER: &'static str = include_str!("vertex_shader.mesa.glsl");
#[cfg(target_os = "linux")]
static FRAGMENT_SHADER: &'static str = include_str!("fragment_shader.mesa.glsl");
<|endoftext|>


<commit_before>use crate::beacon_node_fallback::{AllErrored, Error as FallbackError};
use crate::{
    beacon_node_fallback::{BeaconNodeFallback, RequireSynced},
    graffiti_file::GraffitiFile,
};
use crate::{http_metrics::metrics, validator_store::ValidatorStore};
use environment::RuntimeContext;
use eth2::types::Graffiti;
use slog::{crit, debug, error, info, trace, warn};
use slot_clock::SlotClock;
use std::ops::Deref;
use std::sync::Arc;
use tokio::sync::mpsc;
use types::{
    BlindedPayload, BlockType, Epoch, EthSpec, ExecPayload, FullPayload, PublicKeyBytes, Slot,
};

#[derive(Debug)]
pub enum BlockError {
    Recoverable(String),
    Irrecoverable(String),
}

impl From<AllErrored<BlockError>> for BlockError {
    fn from(e: AllErrored<BlockError>) -> Self {
        if e.0.iter().any(|(_, error)| {
            matches!(
                error,
                FallbackError::RequestFailed(BlockError::Irrecoverable(_))
            )
        }) {
            BlockError::Irrecoverable(e.to_string())
        } else {
            BlockError::Recoverable(e.to_string())
        }
    }
}

/// Builds a `BlockService`.
pub struct BlockServiceBuilder<T, E: EthSpec> {
    validator_store: Option<Arc<ValidatorStore<T, E>>>,
    slot_clock: Option<Arc<T>>,
    beacon_nodes: Option<Arc<BeaconNodeFallback<T, E>>>,
    context: Option<RuntimeContext<E>>,
    graffiti: Option<Graffiti>,
    graffiti_file: Option<GraffitiFile>,
    private_tx_proposals: bool,
}

impl<T: SlotClock + 'static, E: EthSpec> BlockServiceBuilder<T, E> {
    pub fn new() -> Self {
        Self {
            validator_store: None,
            slot_clock: None,
            beacon_nodes: None,
            context: None,
            graffiti: None,
            graffiti_file: None,
            private_tx_proposals: false,
        }
    }

    pub fn validator_store(mut self, store: Arc<ValidatorStore<T, E>>) -> Self {
        self.validator_store = Some(store);
        self
    }

    pub fn slot_clock(mut self, slot_clock: T) -> Self {
        self.slot_clock = Some(Arc::new(slot_clock));
        self
    }

    pub fn beacon_nodes(mut self, beacon_nodes: Arc<BeaconNodeFallback<T, E>>) -> Self {
        self.beacon_nodes = Some(beacon_nodes);
        self
    }

    pub fn runtime_context(mut self, context: RuntimeContext<E>) -> Self {
        self.context = Some(context);
        self
    }

    pub fn graffiti(mut self, graffiti: Option<Graffiti>) -> Self {
        self.graffiti = graffiti;
        self
    }

    pub fn graffiti_file(mut self, graffiti_file: Option<GraffitiFile>) -> Self {
        self.graffiti_file = graffiti_file;
        self
    }

    pub fn private_tx_proposals(mut self, private_tx_proposals: bool) -> Self {
        self.private_tx_proposals = private_tx_proposals;
        self
    }

    pub fn build(self) -> Result<BlockService<T, E>, String> {
        Ok(BlockService {
            inner: Arc::new(Inner {
                validator_store: self
                    .validator_store
                    .ok_or("Cannot build BlockService without validator_store")?,
                slot_clock: self
                    .slot_clock
                    .ok_or("Cannot build BlockService without slot_clock")?,
                beacon_nodes: self
                    .beacon_nodes
                    .ok_or("Cannot build BlockService without beacon_node")?,
                context: self
                    .context
                    .ok_or("Cannot build BlockService without runtime_context")?,
                graffiti: self.graffiti,
                graffiti_file: self.graffiti_file,
                private_tx_proposals: self.private_tx_proposals,
            }),
        })
    }
}

/// Helper to minimise `Arc` usage.
pub struct Inner<T, E: EthSpec> {
    validator_store: Arc<ValidatorStore<T, E>>,
    slot_clock: Arc<T>,
    beacon_nodes: Arc<BeaconNodeFallback<T, E>>,
    context: RuntimeContext<E>,
    graffiti: Option<Graffiti>,
    graffiti_file: Option<GraffitiFile>,
    private_tx_proposals: bool,
}

/// Attempts to produce attestations for any block producer(s) at the start of the epoch.
pub struct BlockService<T, E: EthSpec> {
    inner: Arc<Inner<T, E>>,
}

impl<T, E: EthSpec> Clone for BlockService<T, E> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
        }
    }
}

impl<T, E: EthSpec> Deref for BlockService<T, E> {
    type Target = Inner<T, E>;

    fn deref(&self) -> &Self::Target {
        self.inner.deref()
    }
}

/// Notification from the duties service that we should try to produce a block.
pub struct BlockServiceNotification {
    pub slot: Slot,
    pub block_proposers: Vec<PublicKeyBytes>,
}

impl<T: SlotClock + 'static, E: EthSpec> BlockService<T, E> {
    pub fn start_update_service(
        self,
        mut notification_rx: mpsc::Receiver<BlockServiceNotification>,
    ) -> Result<(), String> {
        let log = self.context.log().clone();

        info!(log, "Block production service started");

        let executor = self.inner.context.executor.clone();

        executor.spawn(
            async move {
                while let Some(notif) = notification_rx.recv().await {
                    let service = self.clone();
                    service.do_update(notif).await.ok();
                }
                debug!(log, "Block service shutting down");
            },
            "block_service",
        );

        Ok(())
    }

    /// Attempt to produce a block for any block producers in the `ValidatorStore`.
    async fn do_update(&self, notification: BlockServiceNotification) -> Result<(), ()> {
        let log = self.context.log();
        let _timer =
            metrics::start_timer_vec(&metrics::BLOCK_SERVICE_TIMES, &[metrics::FULL_UPDATE]);

        let slot = self.slot_clock.now().ok_or_else(move || {
            crit!(log, "Duties manager failed to read slot clock");
        })?;

        if notification.slot != slot {
            warn!(
                log,
                "Skipping block production for expired slot";
                "current_slot" => slot.as_u64(),
                "notification_slot" => notification.slot.as_u64(),
                "info" => "Your machine could be overloaded"
            );
            return Ok(());
        }

        if slot == self.context.eth2_config.spec.genesis_slot {
            debug!(
                log,
                "Not producing block at genesis slot";
                "proposers" => format!("{:?}", notification.block_proposers),
            );
            return Ok(());
        }

        trace!(
            log,
            "Block service update started";
            "slot" => slot.as_u64()
        );

        let proposers = notification.block_proposers;

        if proposers.is_empty() {
            trace!(
                log,
                "No local block proposers for this slot";
                "slot" => slot.as_u64()
            )
        } else if proposers.len() > 1 {
            error!(
                log,
                "Multiple block proposers for this slot";
                "action" => "producing blocks for all proposers",
                "num_proposers" => proposers.len(),
                "slot" => slot.as_u64(),
            )
        }

        let private_tx_proposals = self.private_tx_proposals;
        let merge_slot = self
            .context
            .eth2_config
            .spec
            .bellatrix_fork_epoch
            .unwrap_or_else(Epoch::max_value)
            .start_slot(E::slots_per_epoch());
        for validator_pubkey in proposers {
            let service = self.clone();
            let log = log.clone();
            self.inner.context.executor.spawn(
                async move {
                    let publish_result = if private_tx_proposals && slot >= merge_slot {
                        let mut result = service.clone()
                            .publish_block::<BlindedPayload<E>>(slot, validator_pubkey)
                            .await;
                        match result.as_ref() {
                            Err(BlockError::Recoverable(e)) => {
                                error!(log, "Error whilst producing a blinded block, attempting to publish full block"; "error" => ?e);
                                result = service
                                    .publish_block::<FullPayload<E>>(slot, validator_pubkey)
                                    .await;
                            },
                            Err(BlockError::Irrecoverable(e))  => {
                                error!(log, "Error whilst producing a blinded block, cannot fallback because block was signed"; "error" => ?e);
                            },
                            _ => {},
                        };
                        result
                    } else {
                        service
                            .publish_block::<FullPayload<E>>(slot, validator_pubkey)
                            .await
                    };
                    if let Err(e) = publish_result {
                        crit!(
                            log,
                            "Error whilst producing block";
                            "message" => ?e
                        );
                    }
                },
                "block service",
            );
        }

        Ok(())
    }

    /// Produce a block at the given slot for validator_pubkey
    async fn publish_block<Payload: ExecPayload<E>>(
        self,
        slot: Slot,
        validator_pubkey: PublicKeyBytes,
    ) -> Result<(), BlockError> {
        let log = self.context.log();
        let _timer =
            metrics::start_timer_vec(&metrics::BLOCK_SERVICE_TIMES, &[metrics::BEACON_BLOCK]);

        let current_slot = self.slot_clock.now().ok_or_else(|| {
            BlockError::Recoverable("Unable to determine current slot from clock".to_string())
        })?;

        let randao_reveal = self
            .validator_store
            .randao_reveal(validator_pubkey, slot.epoch(E::slots_per_epoch()))
            .await
            .map_err(|e| {
                BlockError::Recoverable(format!(
                    "Unable to produce randao reveal signature: {:?}",
                    e
                ))
            })?
            .into();

        let graffiti = self
            .graffiti_file
            .clone()
            .and_then(|mut g| match g.load_graffiti(&validator_pubkey) {
                Ok(g) => g,
                Err(e) => {
                    warn!(log, "Failed to read graffiti file"; "error" => ?e);
                    None
                }
            })
            .or_else(|| self.validator_store.graffiti(&validator_pubkey))
            .or(self.graffiti);

        let randao_reveal_ref = &randao_reveal;
        let self_ref = &self;
        let proposer_index = self.validator_store.validator_index(&validator_pubkey);
        let validator_pubkey_ref = &validator_pubkey;
        let signed_block = self
            .beacon_nodes
            .first_success(RequireSynced::No, |beacon_node| async move {
                let get_timer = metrics::start_timer_vec(
                    &metrics::BLOCK_SERVICE_TIMES,
                    &[metrics::BEACON_BLOCK_HTTP_GET],
                );
                let block = match Payload::block_type() {
                    BlockType::Full => {
                        beacon_node
                            .get_validator_blocks::<E, Payload>(
                                slot,
                                randao_reveal_ref,
                                graffiti.as_ref(),
                            )
                            .await
                            .map_err(|e| {
                                BlockError::Recoverable(format!(
                                    "Error from beacon node when producing block: {:?}",
                                    e
                                ))
                            })?
                            .data
                    }
                    BlockType::Blinded => {
                        beacon_node
                            .get_validator_blinded_blocks::<E, Payload>(
                                slot,
                                randao_reveal_ref,
                                graffiti.as_ref(),
                            )
                            .await
                            .map_err(|e| {
                                BlockError::Recoverable(format!(
                                    "Error from beacon node when producing block: {:?}",
                                    e
                                ))
                            })?
                            .data
                    }
                };
                drop(get_timer);

                if proposer_index != Some(block.proposer_index()) {
                    return Err(BlockError::Recoverable(
                        "Proposer index does not match block proposer. Beacon chain re-orged"
                            .to_string(),
                    ));
                }

                let signed_block = self_ref
                    .validator_store
                    .sign_block::<Payload>(*validator_pubkey_ref, block, current_slot)
                    .await
                    .map_err(|e| {
                        BlockError::Recoverable(format!("Unable to sign block: {:?}", e))
                    })?;

                let _post_timer = metrics::start_timer_vec(
                    &metrics::BLOCK_SERVICE_TIMES,
                    &[metrics::BEACON_BLOCK_HTTP_POST],
                );

                match Payload::block_type() {
                    BlockType::Full => beacon_node
                        .post_beacon_blocks(&signed_block)
                        .await
                        .map_err(|e| {
                            BlockError::Irrecoverable(format!(
                                "Error from beacon node when publishing block: {:?}",
                                e
                            ))
                        })?,
                    BlockType::Blinded => beacon_node
                        .post_beacon_blinded_blocks(&signed_block)
                        .await
                        .map_err(|e| {
                            BlockError::Irrecoverable(format!(
                                "Error from beacon node when publishing block: {:?}",
                                e
                            ))
                        })?,
                }

                Ok::<_, BlockError>(signed_block)
            })
            .await?;

        info!(
            log,
            "Successfully published block";
            "deposits" => signed_block.message().body().deposits().len(),
            "attestations" => signed_block.message().body().attestations().len(),
            "graffiti" => ?graffiti.map(|g| g.as_utf8_lossy()),
            "slot" => signed_block.slot().as_u64(),
        );

        Ok(())
    }
}
<commit_msg>Avoid unnecessary slashing protection when publishing blocks (#3188)<commit_after>use crate::beacon_node_fallback::{AllErrored, Error as FallbackError};
use crate::{
    beacon_node_fallback::{BeaconNodeFallback, RequireSynced},
    graffiti_file::GraffitiFile,
};
use crate::{http_metrics::metrics, validator_store::ValidatorStore};
use environment::RuntimeContext;
use eth2::types::Graffiti;
use slog::{crit, debug, error, info, trace, warn};
use slot_clock::SlotClock;
use std::ops::Deref;
use std::sync::Arc;
use tokio::sync::mpsc;
use types::{
    BlindedPayload, BlockType, Epoch, EthSpec, ExecPayload, FullPayload, PublicKeyBytes, Slot,
};

#[derive(Debug)]
pub enum BlockError {
    Recoverable(String),
    Irrecoverable(String),
}

impl From<AllErrored<BlockError>> for BlockError {
    fn from(e: AllErrored<BlockError>) -> Self {
        if e.0.iter().any(|(_, error)| {
            matches!(
                error,
                FallbackError::RequestFailed(BlockError::Irrecoverable(_))
            )
        }) {
            BlockError::Irrecoverable(e.to_string())
        } else {
            BlockError::Recoverable(e.to_string())
        }
    }
}

/// Builds a `BlockService`.
pub struct BlockServiceBuilder<T, E: EthSpec> {
    validator_store: Option<Arc<ValidatorStore<T, E>>>,
    slot_clock: Option<Arc<T>>,
    beacon_nodes: Option<Arc<BeaconNodeFallback<T, E>>>,
    context: Option<RuntimeContext<E>>,
    graffiti: Option<Graffiti>,
    graffiti_file: Option<GraffitiFile>,
    private_tx_proposals: bool,
}

impl<T: SlotClock + 'static, E: EthSpec> BlockServiceBuilder<T, E> {
    pub fn new() -> Self {
        Self {
            validator_store: None,
            slot_clock: None,
            beacon_nodes: None,
            context: None,
            graffiti: None,
            graffiti_file: None,
            private_tx_proposals: false,
        }
    }

    pub fn validator_store(mut self, store: Arc<ValidatorStore<T, E>>) -> Self {
        self.validator_store = Some(store);
        self
    }

    pub fn slot_clock(mut self, slot_clock: T) -> Self {
        self.slot_clock = Some(Arc::new(slot_clock));
        self
    }

    pub fn beacon_nodes(mut self, beacon_nodes: Arc<BeaconNodeFallback<T, E>>) -> Self {
        self.beacon_nodes = Some(beacon_nodes);
        self
    }

    pub fn runtime_context(mut self, context: RuntimeContext<E>) -> Self {
        self.context = Some(context);
        self
    }

    pub fn graffiti(mut self, graffiti: Option<Graffiti>) -> Self {
        self.graffiti = graffiti;
        self
    }

    pub fn graffiti_file(mut self, graffiti_file: Option<GraffitiFile>) -> Self {
        self.graffiti_file = graffiti_file;
        self
    }

    pub fn private_tx_proposals(mut self, private_tx_proposals: bool) -> Self {
        self.private_tx_proposals = private_tx_proposals;
        self
    }

    pub fn build(self) -> Result<BlockService<T, E>, String> {
        Ok(BlockService {
            inner: Arc::new(Inner {
                validator_store: self
                    .validator_store
                    .ok_or("Cannot build BlockService without validator_store")?,
                slot_clock: self
                    .slot_clock
                    .ok_or("Cannot build BlockService without slot_clock")?,
                beacon_nodes: self
                    .beacon_nodes
                    .ok_or("Cannot build BlockService without beacon_node")?,
                context: self
                    .context
                    .ok_or("Cannot build BlockService without runtime_context")?,
                graffiti: self.graffiti,
                graffiti_file: self.graffiti_file,
                private_tx_proposals: self.private_tx_proposals,
            }),
        })
    }
}

/// Helper to minimise `Arc` usage.
pub struct Inner<T, E: EthSpec> {
    validator_store: Arc<ValidatorStore<T, E>>,
    slot_clock: Arc<T>,
    beacon_nodes: Arc<BeaconNodeFallback<T, E>>,
    context: RuntimeContext<E>,
    graffiti: Option<Graffiti>,
    graffiti_file: Option<GraffitiFile>,
    private_tx_proposals: bool,
}

/// Attempts to produce attestations for any block producer(s) at the start of the epoch.
pub struct BlockService<T, E: EthSpec> {
    inner: Arc<Inner<T, E>>,
}

impl<T, E: EthSpec> Clone for BlockService<T, E> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
        }
    }
}

impl<T, E: EthSpec> Deref for BlockService<T, E> {
    type Target = Inner<T, E>;

    fn deref(&self) -> &Self::Target {
        self.inner.deref()
    }
}

/// Notification from the duties service that we should try to produce a block.
pub struct BlockServiceNotification {
    pub slot: Slot,
    pub block_proposers: Vec<PublicKeyBytes>,
}

impl<T: SlotClock + 'static, E: EthSpec> BlockService<T, E> {
    pub fn start_update_service(
        self,
        mut notification_rx: mpsc::Receiver<BlockServiceNotification>,
    ) -> Result<(), String> {
        let log = self.context.log().clone();

        info!(log, "Block production service started");

        let executor = self.inner.context.executor.clone();

        executor.spawn(
            async move {
                while let Some(notif) = notification_rx.recv().await {
                    let service = self.clone();
                    service.do_update(notif).await.ok();
                }
                debug!(log, "Block service shutting down");
            },
            "block_service",
        );

        Ok(())
    }

    /// Attempt to produce a block for any block producers in the `ValidatorStore`.
    async fn do_update(&self, notification: BlockServiceNotification) -> Result<(), ()> {
        let log = self.context.log();
        let _timer =
            metrics::start_timer_vec(&metrics::BLOCK_SERVICE_TIMES, &[metrics::FULL_UPDATE]);

        let slot = self.slot_clock.now().ok_or_else(move || {
            crit!(log, "Duties manager failed to read slot clock");
        })?;

        if notification.slot != slot {
            warn!(
                log,
                "Skipping block production for expired slot";
                "current_slot" => slot.as_u64(),
                "notification_slot" => notification.slot.as_u64(),
                "info" => "Your machine could be overloaded"
            );
            return Ok(());
        }

        if slot == self.context.eth2_config.spec.genesis_slot {
            debug!(
                log,
                "Not producing block at genesis slot";
                "proposers" => format!("{:?}", notification.block_proposers),
            );
            return Ok(());
        }

        trace!(
            log,
            "Block service update started";
            "slot" => slot.as_u64()
        );

        let proposers = notification.block_proposers;

        if proposers.is_empty() {
            trace!(
                log,
                "No local block proposers for this slot";
                "slot" => slot.as_u64()
            )
        } else if proposers.len() > 1 {
            error!(
                log,
                "Multiple block proposers for this slot";
                "action" => "producing blocks for all proposers",
                "num_proposers" => proposers.len(),
                "slot" => slot.as_u64(),
            )
        }

        let private_tx_proposals = self.private_tx_proposals;
        let merge_slot = self
            .context
            .eth2_config
            .spec
            .bellatrix_fork_epoch
            .unwrap_or_else(Epoch::max_value)
            .start_slot(E::slots_per_epoch());
        for validator_pubkey in proposers {
            let service = self.clone();
            let log = log.clone();
            self.inner.context.executor.spawn(
                async move {
                    let publish_result = if private_tx_proposals && slot >= merge_slot {
                        let mut result = service.clone()
                            .publish_block::<BlindedPayload<E>>(slot, validator_pubkey)
                            .await;
                        match result.as_ref() {
                            Err(BlockError::Recoverable(e)) => {
                                error!(log, "Error whilst producing a blinded block, attempting to publish full block"; "error" => ?e);
                                result = service
                                    .publish_block::<FullPayload<E>>(slot, validator_pubkey)
                                    .await;
                            },
                            Err(BlockError::Irrecoverable(e))  => {
                                error!(log, "Error whilst producing a blinded block, cannot fallback because block was signed"; "error" => ?e);
                            },
                            _ => {},
                        };
                        result
                    } else {
                        service
                            .publish_block::<FullPayload<E>>(slot, validator_pubkey)
                            .await
                    };
                    if let Err(e) = publish_result {
                        crit!(
                            log,
                            "Error whilst producing block";
                            "message" => ?e
                        );
                    }
                },
                "block service",
            );
        }

        Ok(())
    }

    /// Produce a block at the given slot for validator_pubkey
    async fn publish_block<Payload: ExecPayload<E>>(
        self,
        slot: Slot,
        validator_pubkey: PublicKeyBytes,
    ) -> Result<(), BlockError> {
        let log = self.context.log();
        let _timer =
            metrics::start_timer_vec(&metrics::BLOCK_SERVICE_TIMES, &[metrics::BEACON_BLOCK]);

        let current_slot = self.slot_clock.now().ok_or_else(|| {
            BlockError::Recoverable("Unable to determine current slot from clock".to_string())
        })?;

        let randao_reveal = self
            .validator_store
            .randao_reveal(validator_pubkey, slot.epoch(E::slots_per_epoch()))
            .await
            .map_err(|e| {
                BlockError::Recoverable(format!(
                    "Unable to produce randao reveal signature: {:?}",
                    e
                ))
            })?
            .into();

        let graffiti = self
            .graffiti_file
            .clone()
            .and_then(|mut g| match g.load_graffiti(&validator_pubkey) {
                Ok(g) => g,
                Err(e) => {
                    warn!(log, "Failed to read graffiti file"; "error" => ?e);
                    None
                }
            })
            .or_else(|| self.validator_store.graffiti(&validator_pubkey))
            .or(self.graffiti);

        let randao_reveal_ref = &randao_reveal;
        let self_ref = &self;
        let proposer_index = self.validator_store.validator_index(&validator_pubkey);
        let validator_pubkey_ref = &validator_pubkey;
        // Request block from first responsive beacon node.
        let block = self
            .beacon_nodes
            .first_success(RequireSynced::No, |beacon_node| async move {
                let get_timer = metrics::start_timer_vec(
                    &metrics::BLOCK_SERVICE_TIMES,
                    &[metrics::BEACON_BLOCK_HTTP_GET],
                );
                let block = match Payload::block_type() {
                    BlockType::Full => {
                        beacon_node
                            .get_validator_blocks::<E, Payload>(
                                slot,
                                randao_reveal_ref,
                                graffiti.as_ref(),
                            )
                            .await
                            .map_err(|e| {
                                BlockError::Recoverable(format!(
                                    "Error from beacon node when producing block: {:?}",
                                    e
                                ))
                            })?
                            .data
                    }
                    BlockType::Blinded => {
                        beacon_node
                            .get_validator_blinded_blocks::<E, Payload>(
                                slot,
                                randao_reveal_ref,
                                graffiti.as_ref(),
                            )
                            .await
                            .map_err(|e| {
                                BlockError::Recoverable(format!(
                                    "Error from beacon node when producing block: {:?}",
                                    e
                                ))
                            })?
                            .data
                    }
                };
                drop(get_timer);

                if proposer_index != Some(block.proposer_index()) {
                    return Err(BlockError::Recoverable(
                        "Proposer index does not match block proposer. Beacon chain re-orged"
                            .to_string(),
                    ));
                }

                Ok::<_, BlockError>(block)
            })
            .await?;

        let signed_block = self_ref
            .validator_store
            .sign_block::<Payload>(*validator_pubkey_ref, block, current_slot)
            .await
            .map_err(|e| BlockError::Recoverable(format!("Unable to sign block: {:?}", e)))?;

        // Publish block with first available beacon node.
        self.beacon_nodes
            .first_success(RequireSynced::No, |beacon_node| async {
                let _post_timer = metrics::start_timer_vec(
                    &metrics::BLOCK_SERVICE_TIMES,
                    &[metrics::BEACON_BLOCK_HTTP_POST],
                );

                match Payload::block_type() {
                    BlockType::Full => beacon_node
                        .post_beacon_blocks(&signed_block)
                        .await
                        .map_err(|e| {
                            BlockError::Irrecoverable(format!(
                                "Error from beacon node when publishing block: {:?}",
                                e
                            ))
                        })?,
                    BlockType::Blinded => beacon_node
                        .post_beacon_blinded_blocks(&signed_block)
                        .await
                        .map_err(|e| {
                            BlockError::Irrecoverable(format!(
                                "Error from beacon node when publishing block: {:?}",
                                e
                            ))
                        })?,
                }

                info!(
                    log,
                    "Successfully published block";
                    "deposits" => signed_block.message().body().deposits().len(),
                    "attestations" => signed_block.message().body().attestations().len(),
                    "graffiti" => ?graffiti.map(|g| g.as_utf8_lossy()),
                    "slot" => signed_block.slot().as_u64(),
                );
                Ok::<_, BlockError>(())
            })
            .await?;
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! To make attribute macros work reliably when typing, we need to take care to
//! fix up syntax errors in the code we're passing to them.
use std::mem;

use mbe::{SyntheticToken, SyntheticTokenId, TokenMap};
use rustc_hash::FxHashMap;
use syntax::{
    ast::{self, AstNode},
    match_ast, SyntaxKind, SyntaxNode, TextRange,
};
use tt::Subtree;

/// The result of calculating fixes for a syntax node -- a bunch of changes
/// (appending to and replacing nodes), the information that is needed to
/// reverse those changes afterwards, and a token map.
#[derive(Debug)]
pub(crate) struct SyntaxFixups {
    pub(crate) append: FxHashMap<SyntaxNode, Vec<SyntheticToken>>,
    pub(crate) replace: FxHashMap<SyntaxNode, Vec<SyntheticToken>>,
    pub(crate) undo_info: SyntaxFixupUndoInfo,
    pub(crate) token_map: TokenMap,
    pub(crate) next_id: u32,
}

/// This is the information needed to reverse the fixups.
#[derive(Debug, PartialEq, Eq)]
pub struct SyntaxFixupUndoInfo {
    original: Vec<Subtree>,
}

const EMPTY_ID: SyntheticTokenId = SyntheticTokenId(!0);

pub(crate) fn fixup_syntax(node: &SyntaxNode) -> SyntaxFixups {
    let mut append = FxHashMap::default();
    let mut replace = FxHashMap::default();
    let mut preorder = node.preorder();
    let mut original = Vec::new();
    let mut token_map = TokenMap::default();
    let mut next_id = 0;
    while let Some(event) = preorder.next() {
        let node = match event {
            syntax::WalkEvent::Enter(node) => node,
            syntax::WalkEvent::Leave(_) => continue,
        };

        if can_handle_error(&node) && has_error_to_handle(&node) {
            // the node contains an error node, we have to completely replace it by something valid
            let (original_tree, new_tmap, new_next_id) =
                mbe::syntax_node_to_token_tree_with_modifications(
                    &node,
                    mem::take(&mut token_map),
                    next_id,
                    Default::default(),
                    Default::default(),
                );
            token_map = new_tmap;
            next_id = new_next_id;
            let idx = original.len() as u32;
            original.push(original_tree);
            let replacement = SyntheticToken {
                kind: SyntaxKind::IDENT,
                text: "__ra_fixup".into(),
                range: node.text_range(),
                id: SyntheticTokenId(idx),
            };
            replace.insert(node.clone(), vec![replacement]);
            preorder.skip_subtree();
            continue;
        }

        // In some other situations, we can fix things by just appending some tokens.
        let end_range = TextRange::empty(node.text_range().end());
        match_ast! {
            match node {
                ast::FieldExpr(it) => {
                    if it.name_ref().is_none() {
                        // incomplete field access: some_expr.|
                        append.insert(node.clone(), vec![
                            SyntheticToken {
                                kind: SyntaxKind::IDENT,
                                text: "__ra_fixup".into(),
                                range: end_range,
                                id: EMPTY_ID,
                            },
                        ]);
                    }
                },
                ast::ExprStmt(it) => {
                    if it.semicolon_token().is_none() {
                        append.insert(node.clone(), vec![
                            SyntheticToken {
                                kind: SyntaxKind::SEMICOLON,
                                text: ";".into(),
                                range: end_range,
                                id: EMPTY_ID,
                            },
                        ]);
                    }
                },
                _ => (),
            }
        }
    }
    SyntaxFixups {
        append,
        replace,
        token_map,
        next_id,
        undo_info: SyntaxFixupUndoInfo { original },
    }
}

fn has_error(node: &SyntaxNode) -> bool {
    node.children().any(|c| c.kind() == SyntaxKind::ERROR)
}

fn can_handle_error(node: &SyntaxNode) -> bool {
    ast::Expr::can_cast(node.kind())
}

fn has_error_to_handle(node: &SyntaxNode) -> bool {
    has_error(node) || node.children().any(|c| !can_handle_error(&c) && has_error_to_handle(&c))
}

pub(crate) fn reverse_fixups(tt: &mut Subtree, token_map: &TokenMap, undo_info: &SyntaxFixupUndoInfo) {
    tt.token_trees.retain(|tt| match tt {
        tt::TokenTree::Leaf(leaf) => {
            token_map.synthetic_token_id(leaf.id()).is_none()
                || token_map.synthetic_token_id(leaf.id()) != Some(EMPTY_ID)
        }
        _ => true,
    });
    tt.token_trees.iter_mut().for_each(|tt| match tt {
        tt::TokenTree::Subtree(tt) => reverse_fixups(tt, token_map, undo_info),
        tt::TokenTree::Leaf(leaf) => {
            if let Some(id) = token_map.synthetic_token_id(leaf.id()) {
                let original = &undo_info.original[id.0 as usize];
                *tt = tt::TokenTree::Subtree(original.clone());
            }
        }
    });
}

#[cfg(test)]
mod tests {
    use expect_test::{expect, Expect};

    use super::reverse_fixups;

    #[track_caller]
    fn check(ra_fixture: &str, mut expect: Expect) {
        let parsed = syntax::SourceFile::parse(ra_fixture);
        let fixups = super::fixup_syntax(&parsed.syntax_node());
        let (mut tt, tmap, _) = mbe::syntax_node_to_token_tree_with_modifications(
            &parsed.syntax_node(),
            fixups.token_map,
            fixups.next_id,
            fixups.replace,
            fixups.append,
        );

        let mut actual = tt.to_string();
        actual.push_str("\n");

        expect.indent(false);
        expect.assert_eq(&actual);

        // the fixed-up tree should be syntactically valid
        let (parse, _) = mbe::token_tree_to_syntax_node(&tt, ::mbe::TopEntryPoint::MacroItems);
        assert_eq!(
            parse.errors(),
            &[],
            "parse has syntax errors. parse tree:\n{:#?}",
            parse.syntax_node()
        );

        reverse_fixups(&mut tt, &tmap, &fixups.undo_info);

        // the fixed-up + reversed version should be equivalent to the original input
        // (but token IDs don't matter)
        let (original_as_tt, _) = mbe::syntax_node_to_token_tree(&parsed.syntax_node());
        assert_eq!(tt.to_string(), original_as_tt.to_string());
    }

    #[test]
    fn incomplete_field_expr_1() {
        check(
            r#"
fn foo() {
    a.
}
"#,
            expect![[r#"
fn foo () {a . __ra_fixup}
"#]],
        )
    }

    #[test]
    fn incomplete_field_expr_2() {
        check(
            r#"
fn foo() {
    a. ;
}
"#,
            expect![[r#"
fn foo () {a . __ra_fixup ;}
"#]],
        )
    }

    #[test]
    fn incomplete_field_expr_3() {
        check(
            r#"
fn foo() {
    a. ;
    bar();
}
"#,
            expect![[r#"
fn foo () {a . __ra_fixup ; bar () ;}
"#]],
        )
    }

    #[test]
    fn field_expr_before_call() {
        // another case that easily happens while typing
        check(
            r#"
fn foo() {
    a.b
    bar();
}
"#,
            expect![[r#"
fn foo () {a . b ; bar () ;}
"#]],
        )
    }

    #[test]
    fn extraneous_comma() {
        check(
            r#"
fn foo() {
    bar(,);
}
"#,
            expect![[r#"
fn foo () {__ra_fixup ;}
"#]],
        )
    }
}
<commit_msg>Format again<commit_after>//! To make attribute macros work reliably when typing, we need to take care to
//! fix up syntax errors in the code we're passing to them.
use std::mem;

use mbe::{SyntheticToken, SyntheticTokenId, TokenMap};
use rustc_hash::FxHashMap;
use syntax::{
    ast::{self, AstNode},
    match_ast, SyntaxKind, SyntaxNode, TextRange,
};
use tt::Subtree;

/// The result of calculating fixes for a syntax node -- a bunch of changes
/// (appending to and replacing nodes), the information that is needed to
/// reverse those changes afterwards, and a token map.
#[derive(Debug)]
pub(crate) struct SyntaxFixups {
    pub(crate) append: FxHashMap<SyntaxNode, Vec<SyntheticToken>>,
    pub(crate) replace: FxHashMap<SyntaxNode, Vec<SyntheticToken>>,
    pub(crate) undo_info: SyntaxFixupUndoInfo,
    pub(crate) token_map: TokenMap,
    pub(crate) next_id: u32,
}

/// This is the information needed to reverse the fixups.
#[derive(Debug, PartialEq, Eq)]
pub struct SyntaxFixupUndoInfo {
    original: Vec<Subtree>,
}

const EMPTY_ID: SyntheticTokenId = SyntheticTokenId(!0);

pub(crate) fn fixup_syntax(node: &SyntaxNode) -> SyntaxFixups {
    let mut append = FxHashMap::default();
    let mut replace = FxHashMap::default();
    let mut preorder = node.preorder();
    let mut original = Vec::new();
    let mut token_map = TokenMap::default();
    let mut next_id = 0;
    while let Some(event) = preorder.next() {
        let node = match event {
            syntax::WalkEvent::Enter(node) => node,
            syntax::WalkEvent::Leave(_) => continue,
        };

        if can_handle_error(&node) && has_error_to_handle(&node) {
            // the node contains an error node, we have to completely replace it by something valid
            let (original_tree, new_tmap, new_next_id) =
                mbe::syntax_node_to_token_tree_with_modifications(
                    &node,
                    mem::take(&mut token_map),
                    next_id,
                    Default::default(),
                    Default::default(),
                );
            token_map = new_tmap;
            next_id = new_next_id;
            let idx = original.len() as u32;
            original.push(original_tree);
            let replacement = SyntheticToken {
                kind: SyntaxKind::IDENT,
                text: "__ra_fixup".into(),
                range: node.text_range(),
                id: SyntheticTokenId(idx),
            };
            replace.insert(node.clone(), vec![replacement]);
            preorder.skip_subtree();
            continue;
        }

        // In some other situations, we can fix things by just appending some tokens.
        let end_range = TextRange::empty(node.text_range().end());
        match_ast! {
            match node {
                ast::FieldExpr(it) => {
                    if it.name_ref().is_none() {
                        // incomplete field access: some_expr.|
                        append.insert(node.clone(), vec![
                            SyntheticToken {
                                kind: SyntaxKind::IDENT,
                                text: "__ra_fixup".into(),
                                range: end_range,
                                id: EMPTY_ID,
                            },
                        ]);
                    }
                },
                ast::ExprStmt(it) => {
                    if it.semicolon_token().is_none() {
                        append.insert(node.clone(), vec![
                            SyntheticToken {
                                kind: SyntaxKind::SEMICOLON,
                                text: ";".into(),
                                range: end_range,
                                id: EMPTY_ID,
                            },
                        ]);
                    }
                },
                _ => (),
            }
        }
    }
    SyntaxFixups {
        append,
        replace,
        token_map,
        next_id,
        undo_info: SyntaxFixupUndoInfo { original },
    }
}

fn has_error(node: &SyntaxNode) -> bool {
    node.children().any(|c| c.kind() == SyntaxKind::ERROR)
}

fn can_handle_error(node: &SyntaxNode) -> bool {
    ast::Expr::can_cast(node.kind())
}

fn has_error_to_handle(node: &SyntaxNode) -> bool {
    has_error(node) || node.children().any(|c| !can_handle_error(&c) && has_error_to_handle(&c))
}

pub(crate) fn reverse_fixups(
    tt: &mut Subtree,
    token_map: &TokenMap,
    undo_info: &SyntaxFixupUndoInfo,
) {
    tt.token_trees.retain(|tt| match tt {
        tt::TokenTree::Leaf(leaf) => {
            token_map.synthetic_token_id(leaf.id()).is_none()
                || token_map.synthetic_token_id(leaf.id()) != Some(EMPTY_ID)
        }
        _ => true,
    });
    tt.token_trees.iter_mut().for_each(|tt| match tt {
        tt::TokenTree::Subtree(tt) => reverse_fixups(tt, token_map, undo_info),
        tt::TokenTree::Leaf(leaf) => {
            if let Some(id) = token_map.synthetic_token_id(leaf.id()) {
                let original = &undo_info.original[id.0 as usize];
                *tt = tt::TokenTree::Subtree(original.clone());
            }
        }
    });
}

#[cfg(test)]
mod tests {
    use expect_test::{expect, Expect};

    use super::reverse_fixups;

    #[track_caller]
    fn check(ra_fixture: &str, mut expect: Expect) {
        let parsed = syntax::SourceFile::parse(ra_fixture);
        let fixups = super::fixup_syntax(&parsed.syntax_node());
        let (mut tt, tmap, _) = mbe::syntax_node_to_token_tree_with_modifications(
            &parsed.syntax_node(),
            fixups.token_map,
            fixups.next_id,
            fixups.replace,
            fixups.append,
        );

        let mut actual = tt.to_string();
        actual.push_str("\n");

        expect.indent(false);
        expect.assert_eq(&actual);

        // the fixed-up tree should be syntactically valid
        let (parse, _) = mbe::token_tree_to_syntax_node(&tt, ::mbe::TopEntryPoint::MacroItems);
        assert_eq!(
            parse.errors(),
            &[],
            "parse has syntax errors. parse tree:\n{:#?}",
            parse.syntax_node()
        );

        reverse_fixups(&mut tt, &tmap, &fixups.undo_info);

        // the fixed-up + reversed version should be equivalent to the original input
        // (but token IDs don't matter)
        let (original_as_tt, _) = mbe::syntax_node_to_token_tree(&parsed.syntax_node());
        assert_eq!(tt.to_string(), original_as_tt.to_string());
    }

    #[test]
    fn incomplete_field_expr_1() {
        check(
            r#"
fn foo() {
    a.
}
"#,
            expect![[r#"
fn foo () {a . __ra_fixup}
"#]],
        )
    }

    #[test]
    fn incomplete_field_expr_2() {
        check(
            r#"
fn foo() {
    a. ;
}
"#,
            expect![[r#"
fn foo () {a . __ra_fixup ;}
"#]],
        )
    }

    #[test]
    fn incomplete_field_expr_3() {
        check(
            r#"
fn foo() {
    a. ;
    bar();
}
"#,
            expect![[r#"
fn foo () {a . __ra_fixup ; bar () ;}
"#]],
        )
    }

    #[test]
    fn field_expr_before_call() {
        // another case that easily happens while typing
        check(
            r#"
fn foo() {
    a.b
    bar();
}
"#,
            expect![[r#"
fn foo () {a . b ; bar () ;}
"#]],
        )
    }

    #[test]
    fn extraneous_comma() {
        check(
            r#"
fn foo() {
    bar(,);
}
"#,
            expect![[r#"
fn foo () {__ra_fixup ;}
"#]],
        )
    }
}
<|endoftext|>


<commit_before>//! Defines database & queries for macro expansion.

use std::sync::Arc;

use mbe::MacroRules;
use ra_db::{salsa, SourceDatabase};
use ra_parser::FragmentKind;
use ra_prof::profile;
use ra_syntax::{AstNode, Parse, SyntaxKind::*, SyntaxNode};

use crate::{
    ast_id_map::AstIdMap, BuiltinDeriveExpander, BuiltinFnLikeExpander, EagerCallLoc, EagerMacroId,
    HirFileId, HirFileIdRepr, LazyMacroId, MacroCallId, MacroCallLoc, MacroDefId, MacroDefKind,
    MacroFile,
};

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum TokenExpander {
    MacroRules(mbe::MacroRules),
    Builtin(BuiltinFnLikeExpander),
    BuiltinDerive(BuiltinDeriveExpander),
}

impl TokenExpander {
    pub fn expand(
        &self,
        db: &dyn AstDatabase,
        id: LazyMacroId,
        tt: &tt::Subtree,
    ) -> Result<tt::Subtree, mbe::ExpandError> {
        match self {
            TokenExpander::MacroRules(it) => it.expand(tt),
            TokenExpander::Builtin(it) => it.expand(db, id, tt),
            TokenExpander::BuiltinDerive(it) => it.expand(db, id, tt),
        }
    }

    pub fn map_id_down(&self, id: tt::TokenId) -> tt::TokenId {
        match self {
            TokenExpander::MacroRules(it) => it.map_id_down(id),
            TokenExpander::Builtin(..) => id,
            TokenExpander::BuiltinDerive(..) => id,
        }
    }

    pub fn map_id_up(&self, id: tt::TokenId) -> (tt::TokenId, mbe::Origin) {
        match self {
            TokenExpander::MacroRules(it) => it.map_id_up(id),
            TokenExpander::Builtin(..) => (id, mbe::Origin::Call),
            TokenExpander::BuiltinDerive(..) => (id, mbe::Origin::Call),
        }
    }
}

// FIXME: rename to ExpandDatabase
#[salsa::query_group(AstDatabaseStorage)]
pub trait AstDatabase: SourceDatabase {
    fn ast_id_map(&self, file_id: HirFileId) -> Arc<AstIdMap>;

    #[salsa::transparent]
    fn parse_or_expand(&self, file_id: HirFileId) -> Option<SyntaxNode>;

    #[salsa::interned]
    fn intern_macro(&self, macro_call: MacroCallLoc) -> LazyMacroId;
    fn macro_arg(&self, id: MacroCallId) -> Option<Arc<(tt::Subtree, mbe::TokenMap)>>;
    fn macro_def(&self, id: MacroDefId) -> Option<Arc<(TokenExpander, mbe::TokenMap)>>;
    fn parse_macro(&self, macro_file: MacroFile)
        -> Option<(Parse<SyntaxNode>, Arc<mbe::TokenMap>)>;
    fn macro_expand(&self, macro_call: MacroCallId) -> Result<Arc<tt::Subtree>, String>;

    #[salsa::interned]
    fn intern_eager_expansion(&self, eager: EagerCallLoc) -> EagerMacroId;
}

pub(crate) fn ast_id_map(db: &dyn AstDatabase, file_id: HirFileId) -> Arc<AstIdMap> {
    let map =
        db.parse_or_expand(file_id).map_or_else(AstIdMap::default, |it| AstIdMap::from_source(&it));
    Arc::new(map)
}

pub(crate) fn macro_def(
    db: &dyn AstDatabase,
    id: MacroDefId,
) -> Option<Arc<(TokenExpander, mbe::TokenMap)>> {
    match id.kind {
        MacroDefKind::Declarative => {
            let macro_call = id.ast_id?.to_node(db);
            let arg = macro_call.token_tree()?;
            let (tt, tmap) = mbe::ast_to_token_tree(&arg).or_else(|| {
                log::warn!("fail on macro_def to token tree: {:#?}", arg);
                None
            })?;
            let rules = match MacroRules::parse(&tt) {
                Ok(it) => it,
                Err(err) => {
                    log::warn!("fail on macro_def parse: error: {:#?} {:#?}", err, tt);
                    return None;
                }
            };
            Some(Arc::new((TokenExpander::MacroRules(rules), tmap)))
        }
        MacroDefKind::BuiltIn(expander) => {
            Some(Arc::new((TokenExpander::Builtin(expander), mbe::TokenMap::default())))
        }
        MacroDefKind::BuiltInDerive(expander) => {
            Some(Arc::new((TokenExpander::BuiltinDerive(expander), mbe::TokenMap::default())))
        }
        MacroDefKind::BuiltInEager(_expander) => None,
    }
}

pub(crate) fn macro_arg(
    db: &dyn AstDatabase,
    id: MacroCallId,
) -> Option<Arc<(tt::Subtree, mbe::TokenMap)>> {
    let id = match id {
        MacroCallId::LazyMacro(id) => id,
        MacroCallId::EagerMacro(_id) => {
            // FIXME: support macro_arg for eager macro
            return None;
        }
    };
    let loc = db.lookup_intern_macro(id);
    let arg = loc.kind.arg(db)?;
    let (tt, tmap) = mbe::syntax_node_to_token_tree(&arg)?;
    Some(Arc::new((tt, tmap)))
}

pub(crate) fn macro_expand(
    db: &dyn AstDatabase,
    id: MacroCallId,
) -> Result<Arc<tt::Subtree>, String> {
    let lazy_id = match id {
        MacroCallId::LazyMacro(id) => id,
        MacroCallId::EagerMacro(id) => {
            return Ok(db.lookup_intern_eager_expansion(id).subtree);
        }
    };

    let loc = db.lookup_intern_macro(lazy_id);
    let macro_arg = db.macro_arg(id).ok_or("Fail to args in to tt::TokenTree")?;

    let macro_rules = db.macro_def(loc.def).ok_or("Fail to find macro definition")?;
    let tt = macro_rules.0.expand(db, lazy_id, &macro_arg.0).map_err(|err| format!("{:?}", err))?;
    // Set a hard limit for the expanded tt
    let count = tt.count();
    if count > 65536 {
        return Err(format!("Total tokens count exceed limit : count = {}", count));
    }
    Ok(Arc::new(tt))
}

pub(crate) fn parse_or_expand(db: &dyn AstDatabase, file_id: HirFileId) -> Option<SyntaxNode> {
    match file_id.0 {
        HirFileIdRepr::FileId(file_id) => Some(db.parse(file_id).tree().syntax().clone()),
        HirFileIdRepr::MacroFile(macro_file) => {
            db.parse_macro(macro_file).map(|(it, _)| it.syntax_node())
        }
    }
}

pub(crate) fn parse_macro(
    db: &dyn AstDatabase,
    macro_file: MacroFile,
) -> Option<(Parse<SyntaxNode>, Arc<mbe::TokenMap>)> {
    let _p = profile("parse_macro_query");

    let macro_call_id = macro_file.macro_call_id;
    let tt = db
        .macro_expand(macro_call_id)
        .map_err(|err| {
            // Note:
            // The final goal we would like to make all parse_macro success,
            // such that the following log will not call anyway.
            match macro_call_id {
                MacroCallId::LazyMacro(id) => {
                    let loc: MacroCallLoc = db.lookup_intern_macro(id);
                    let node = loc.kind.node(db);
                    log::warn!(
                        "fail on macro_parse: (reason: {} macro_call: {:#})",
                        err,
                        node.value
                    );
                    let mut parent = loc.kind.file_id().call_node(db);
                    while let Some(node) = parent.clone() {
                        log::warn!("parent: macro_call: {:#})", node.value);
                        parent = node.file_id.call_node(db);
                    }
                }
                _ => {
                    log::warn!("fail on macro_parse: (reason: {})", err);
                }
            }
        })
        .ok()?;

    let fragment_kind = to_fragment_kind(db, macro_call_id);

    let (parse, rev_token_map) = mbe::token_tree_to_syntax_node(&tt, fragment_kind).ok()?;
    Some((parse, Arc::new(rev_token_map)))
}

/// Given a `MacroCallId`, return what `FragmentKind` it belongs to.
/// FIXME: Not completed
fn to_fragment_kind(db: &dyn AstDatabase, id: MacroCallId) -> FragmentKind {
    let lazy_id = match id {
        MacroCallId::LazyMacro(id) => id,
        MacroCallId::EagerMacro(id) => {
            return db.lookup_intern_eager_expansion(id).fragment;
        }
    };
    let syn = db.lookup_intern_macro(lazy_id).kind.node(db).value;

    let parent = match syn.parent() {
        Some(it) => it,
        None => {
            // FIXME:
            // If it is root, which means the parent HirFile
            // MacroKindFile must be non-items
            // return expr now.
            return FragmentKind::Expr;
        }
    };

    match parent.kind() {
        MACRO_ITEMS | SOURCE_FILE => FragmentKind::Items,
        ITEM_LIST => FragmentKind::Items,
        LET_STMT => {
            // FIXME: Handle Pattern
            FragmentKind::Expr
        }
        // FIXME: Expand to statements in appropriate positions; HIR lowering needs to handle that
        EXPR_STMT | BLOCK => FragmentKind::Expr,
        ARG_LIST => FragmentKind::Expr,
        TRY_EXPR => FragmentKind::Expr,
        TUPLE_EXPR => FragmentKind::Expr,
        PAREN_EXPR => FragmentKind::Expr,

        FOR_EXPR => FragmentKind::Expr,
        PATH_EXPR => FragmentKind::Expr,
        LAMBDA_EXPR => FragmentKind::Expr,
        CONDITION => FragmentKind::Expr,
        BREAK_EXPR => FragmentKind::Expr,
        RETURN_EXPR => FragmentKind::Expr,
        BLOCK_EXPR => FragmentKind::Expr,
        MATCH_EXPR => FragmentKind::Expr,
        MATCH_ARM => FragmentKind::Expr,
        MATCH_GUARD => FragmentKind::Expr,
        RECORD_FIELD => FragmentKind::Expr,
        CALL_EXPR => FragmentKind::Expr,
        INDEX_EXPR => FragmentKind::Expr,
        METHOD_CALL_EXPR => FragmentKind::Expr,
        AWAIT_EXPR => FragmentKind::Expr,
        CAST_EXPR => FragmentKind::Expr,
        REF_EXPR => FragmentKind::Expr,
        PREFIX_EXPR => FragmentKind::Expr,
        RANGE_EXPR => FragmentKind::Expr,
        BIN_EXPR => FragmentKind::Expr,
        _ => {
            // Unknown , Just guess it is `Items`
            FragmentKind::Items
        }
    }
}
<commit_msg>Add comment for parents logging<commit_after>//! Defines database & queries for macro expansion.

use std::sync::Arc;

use mbe::MacroRules;
use ra_db::{salsa, SourceDatabase};
use ra_parser::FragmentKind;
use ra_prof::profile;
use ra_syntax::{AstNode, Parse, SyntaxKind::*, SyntaxNode};

use crate::{
    ast_id_map::AstIdMap, BuiltinDeriveExpander, BuiltinFnLikeExpander, EagerCallLoc, EagerMacroId,
    HirFileId, HirFileIdRepr, LazyMacroId, MacroCallId, MacroCallLoc, MacroDefId, MacroDefKind,
    MacroFile,
};

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum TokenExpander {
    MacroRules(mbe::MacroRules),
    Builtin(BuiltinFnLikeExpander),
    BuiltinDerive(BuiltinDeriveExpander),
}

impl TokenExpander {
    pub fn expand(
        &self,
        db: &dyn AstDatabase,
        id: LazyMacroId,
        tt: &tt::Subtree,
    ) -> Result<tt::Subtree, mbe::ExpandError> {
        match self {
            TokenExpander::MacroRules(it) => it.expand(tt),
            TokenExpander::Builtin(it) => it.expand(db, id, tt),
            TokenExpander::BuiltinDerive(it) => it.expand(db, id, tt),
        }
    }

    pub fn map_id_down(&self, id: tt::TokenId) -> tt::TokenId {
        match self {
            TokenExpander::MacroRules(it) => it.map_id_down(id),
            TokenExpander::Builtin(..) => id,
            TokenExpander::BuiltinDerive(..) => id,
        }
    }

    pub fn map_id_up(&self, id: tt::TokenId) -> (tt::TokenId, mbe::Origin) {
        match self {
            TokenExpander::MacroRules(it) => it.map_id_up(id),
            TokenExpander::Builtin(..) => (id, mbe::Origin::Call),
            TokenExpander::BuiltinDerive(..) => (id, mbe::Origin::Call),
        }
    }
}

// FIXME: rename to ExpandDatabase
#[salsa::query_group(AstDatabaseStorage)]
pub trait AstDatabase: SourceDatabase {
    fn ast_id_map(&self, file_id: HirFileId) -> Arc<AstIdMap>;

    #[salsa::transparent]
    fn parse_or_expand(&self, file_id: HirFileId) -> Option<SyntaxNode>;

    #[salsa::interned]
    fn intern_macro(&self, macro_call: MacroCallLoc) -> LazyMacroId;
    fn macro_arg(&self, id: MacroCallId) -> Option<Arc<(tt::Subtree, mbe::TokenMap)>>;
    fn macro_def(&self, id: MacroDefId) -> Option<Arc<(TokenExpander, mbe::TokenMap)>>;
    fn parse_macro(&self, macro_file: MacroFile)
        -> Option<(Parse<SyntaxNode>, Arc<mbe::TokenMap>)>;
    fn macro_expand(&self, macro_call: MacroCallId) -> Result<Arc<tt::Subtree>, String>;

    #[salsa::interned]
    fn intern_eager_expansion(&self, eager: EagerCallLoc) -> EagerMacroId;
}

pub(crate) fn ast_id_map(db: &dyn AstDatabase, file_id: HirFileId) -> Arc<AstIdMap> {
    let map =
        db.parse_or_expand(file_id).map_or_else(AstIdMap::default, |it| AstIdMap::from_source(&it));
    Arc::new(map)
}

pub(crate) fn macro_def(
    db: &dyn AstDatabase,
    id: MacroDefId,
) -> Option<Arc<(TokenExpander, mbe::TokenMap)>> {
    match id.kind {
        MacroDefKind::Declarative => {
            let macro_call = id.ast_id?.to_node(db);
            let arg = macro_call.token_tree()?;
            let (tt, tmap) = mbe::ast_to_token_tree(&arg).or_else(|| {
                log::warn!("fail on macro_def to token tree: {:#?}", arg);
                None
            })?;
            let rules = match MacroRules::parse(&tt) {
                Ok(it) => it,
                Err(err) => {
                    log::warn!("fail on macro_def parse: error: {:#?} {:#?}", err, tt);
                    return None;
                }
            };
            Some(Arc::new((TokenExpander::MacroRules(rules), tmap)))
        }
        MacroDefKind::BuiltIn(expander) => {
            Some(Arc::new((TokenExpander::Builtin(expander), mbe::TokenMap::default())))
        }
        MacroDefKind::BuiltInDerive(expander) => {
            Some(Arc::new((TokenExpander::BuiltinDerive(expander), mbe::TokenMap::default())))
        }
        MacroDefKind::BuiltInEager(_expander) => None,
    }
}

pub(crate) fn macro_arg(
    db: &dyn AstDatabase,
    id: MacroCallId,
) -> Option<Arc<(tt::Subtree, mbe::TokenMap)>> {
    let id = match id {
        MacroCallId::LazyMacro(id) => id,
        MacroCallId::EagerMacro(_id) => {
            // FIXME: support macro_arg for eager macro
            return None;
        }
    };
    let loc = db.lookup_intern_macro(id);
    let arg = loc.kind.arg(db)?;
    let (tt, tmap) = mbe::syntax_node_to_token_tree(&arg)?;
    Some(Arc::new((tt, tmap)))
}

pub(crate) fn macro_expand(
    db: &dyn AstDatabase,
    id: MacroCallId,
) -> Result<Arc<tt::Subtree>, String> {
    let lazy_id = match id {
        MacroCallId::LazyMacro(id) => id,
        MacroCallId::EagerMacro(id) => {
            return Ok(db.lookup_intern_eager_expansion(id).subtree);
        }
    };

    let loc = db.lookup_intern_macro(lazy_id);
    let macro_arg = db.macro_arg(id).ok_or("Fail to args in to tt::TokenTree")?;

    let macro_rules = db.macro_def(loc.def).ok_or("Fail to find macro definition")?;
    let tt = macro_rules.0.expand(db, lazy_id, &macro_arg.0).map_err(|err| format!("{:?}", err))?;
    // Set a hard limit for the expanded tt
    let count = tt.count();
    if count > 65536 {
        return Err(format!("Total tokens count exceed limit : count = {}", count));
    }
    Ok(Arc::new(tt))
}

pub(crate) fn parse_or_expand(db: &dyn AstDatabase, file_id: HirFileId) -> Option<SyntaxNode> {
    match file_id.0 {
        HirFileIdRepr::FileId(file_id) => Some(db.parse(file_id).tree().syntax().clone()),
        HirFileIdRepr::MacroFile(macro_file) => {
            db.parse_macro(macro_file).map(|(it, _)| it.syntax_node())
        }
    }
}

pub(crate) fn parse_macro(
    db: &dyn AstDatabase,
    macro_file: MacroFile,
) -> Option<(Parse<SyntaxNode>, Arc<mbe::TokenMap>)> {
    let _p = profile("parse_macro_query");

    let macro_call_id = macro_file.macro_call_id;
    let tt = db
        .macro_expand(macro_call_id)
        .map_err(|err| {
            // Note:
            // The final goal we would like to make all parse_macro success,
            // such that the following log will not call anyway.
            match macro_call_id {
                MacroCallId::LazyMacro(id) => {
                    let loc: MacroCallLoc = db.lookup_intern_macro(id);
                    let node = loc.kind.node(db);

                    // collect parent information for warning log
                    let parents = std::iter::successors(loc.kind.file_id().call_node(db), |it| {
                        it.file_id.call_node(db)
                    })
                    .map(|n| format!("{:#}", n.value))
                    .collect::<Vec<_>>()
                    .join("\n");

                    log::warn!(
                        "fail on macro_parse: (reason: {} macro_call: {:#}) parents: {}",
                        err,
                        node.value,
                        parents
                    );
                }
                _ => {
                    log::warn!("fail on macro_parse: (reason: {})", err);
                }
            }
        })
        .ok()?;

    let fragment_kind = to_fragment_kind(db, macro_call_id);

    let (parse, rev_token_map) = mbe::token_tree_to_syntax_node(&tt, fragment_kind).ok()?;
    Some((parse, Arc::new(rev_token_map)))
}

/// Given a `MacroCallId`, return what `FragmentKind` it belongs to.
/// FIXME: Not completed
fn to_fragment_kind(db: &dyn AstDatabase, id: MacroCallId) -> FragmentKind {
    let lazy_id = match id {
        MacroCallId::LazyMacro(id) => id,
        MacroCallId::EagerMacro(id) => {
            return db.lookup_intern_eager_expansion(id).fragment;
        }
    };
    let syn = db.lookup_intern_macro(lazy_id).kind.node(db).value;

    let parent = match syn.parent() {
        Some(it) => it,
        None => {
            // FIXME:
            // If it is root, which means the parent HirFile
            // MacroKindFile must be non-items
            // return expr now.
            return FragmentKind::Expr;
        }
    };

    match parent.kind() {
        MACRO_ITEMS | SOURCE_FILE => FragmentKind::Items,
        ITEM_LIST => FragmentKind::Items,
        LET_STMT => {
            // FIXME: Handle Pattern
            FragmentKind::Expr
        }
        // FIXME: Expand to statements in appropriate positions; HIR lowering needs to handle that
        EXPR_STMT | BLOCK => FragmentKind::Expr,
        ARG_LIST => FragmentKind::Expr,
        TRY_EXPR => FragmentKind::Expr,
        TUPLE_EXPR => FragmentKind::Expr,
        PAREN_EXPR => FragmentKind::Expr,

        FOR_EXPR => FragmentKind::Expr,
        PATH_EXPR => FragmentKind::Expr,
        LAMBDA_EXPR => FragmentKind::Expr,
        CONDITION => FragmentKind::Expr,
        BREAK_EXPR => FragmentKind::Expr,
        RETURN_EXPR => FragmentKind::Expr,
        BLOCK_EXPR => FragmentKind::Expr,
        MATCH_EXPR => FragmentKind::Expr,
        MATCH_ARM => FragmentKind::Expr,
        MATCH_GUARD => FragmentKind::Expr,
        RECORD_FIELD => FragmentKind::Expr,
        CALL_EXPR => FragmentKind::Expr,
        INDEX_EXPR => FragmentKind::Expr,
        METHOD_CALL_EXPR => FragmentKind::Expr,
        AWAIT_EXPR => FragmentKind::Expr,
        CAST_EXPR => FragmentKind::Expr,
        REF_EXPR => FragmentKind::Expr,
        PREFIX_EXPR => FragmentKind::Expr,
        RANGE_EXPR => FragmentKind::Expr,
        BIN_EXPR => FragmentKind::Expr,
        _ => {
            // Unknown , Just guess it is `Items`
            FragmentKind::Items
        }
    }
}
<|endoftext|>


<commit_before>use std::{
    convert::TryInto,
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use js_int::{uint, UInt};
use serde::{Deserialize, Serialize};

/// A timestamp represented as the number of milliseconds since the unix epoch.
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize)]
#[allow(clippy::exhaustive_structs)]
#[serde(transparent)]
pub struct MilliSecondsSinceUnixEpoch(pub UInt);

impl MilliSecondsSinceUnixEpoch {
    /// Creates a new `MilliSecondsSinceUnixEpoch` from the given `SystemTime`, if it is not before
    /// the unix epoch, or too large to be represented.
    pub fn from_system_time(time: SystemTime) -> Option<Self> {
        let duration = time.duration_since(UNIX_EPOCH).ok()?;
        let millis = duration.as_millis().try_into().ok()?;
        Some(Self(millis))
    }

    /// The current system time in milliseconds since the unix epoch.
    pub fn now() -> Self {
        Self::from_system_time(SystemTime::now()).unwrap()
    }

    /// Creates a new `SystemTime` from `self`, if it can be represented.
    pub fn to_system_time(self) -> Option<SystemTime> {
        UNIX_EPOCH.checked_add(Duration::from_millis(self.0.into()))
    }

    /// Get the time since the unix epoch in milliseconds.
    pub fn get(&self) -> UInt {
        self.0
    }

    /// Get time since the unix epoch in seconds.
    pub fn as_secs(&self) -> UInt {
        self.0 / uint!(1000)
    }
}

/// A timestamp represented as the number of seconds since the unix epoch.
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize)]
#[allow(clippy::exhaustive_structs)]
#[serde(transparent)]
pub struct SecondsSinceUnixEpoch(pub UInt);

impl SecondsSinceUnixEpoch {
    /// Creates a new `MilliSecondsSinceUnixEpoch` from the given `SystemTime`, if it is not before
    /// the unix epoch, or too large to be represented.
    pub fn from_system_time(time: SystemTime) -> Option<Self> {
        let duration = time.duration_since(UNIX_EPOCH).ok()?;
        let millis = duration.as_secs().try_into().ok()?;
        Some(Self(millis))
    }

    /// Creates a new `SystemTime` from `self`, if it can be represented.
    pub fn to_system_time(self) -> Option<SystemTime> {
        UNIX_EPOCH.checked_add(Duration::from_secs(self.0.into()))
    }

    /// Get time since the unix epoch in seconds.
    pub fn get(&self) -> UInt {
        self.0
    }
}

#[cfg(test)]
mod tests {
    use std::time::{Duration, UNIX_EPOCH};

    use js_int::uint;
    use matches::assert_matches;
    use serde::{Deserialize, Serialize};
    use serde_json::json;

    use super::{MilliSecondsSinceUnixEpoch, SecondsSinceUnixEpoch};

    #[derive(Clone, Debug, Deserialize, Serialize)]
    struct SystemTimeTest {
        millis: MilliSecondsSinceUnixEpoch,
        secs: SecondsSinceUnixEpoch,
    }

    #[test]
    fn deserialize() {
        let json = json!({ "millis": 3000, "secs": 60 });

        assert_matches!(
            serde_json::from_value::<SystemTimeTest>(json),
            Ok(SystemTimeTest { millis, secs })
            if millis.to_system_time() == Some(UNIX_EPOCH + Duration::from_millis(3000))
                && secs.to_system_time() == Some(UNIX_EPOCH + Duration::from_secs(60))
        );
    }

    #[test]
    fn serialize() {
        let request = SystemTimeTest {
            millis: MilliSecondsSinceUnixEpoch::from_system_time(UNIX_EPOCH + Duration::new(2, 0))
                .unwrap(),
            secs: SecondsSinceUnixEpoch(uint!(0)),
        };
        assert_matches!(
            serde_json::to_value(&request),
            Ok(value) if value == json!({ "millis": 2000, "secs": 0 })
        );
    }
}
<commit_msg>common: Add SecondsSinceUnixEpoch::now<commit_after>use std::{
    convert::TryInto,
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use js_int::{uint, UInt};
use serde::{Deserialize, Serialize};

/// A timestamp represented as the number of milliseconds since the unix epoch.
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize)]
#[allow(clippy::exhaustive_structs)]
#[serde(transparent)]
pub struct MilliSecondsSinceUnixEpoch(pub UInt);

impl MilliSecondsSinceUnixEpoch {
    /// Creates a new `MilliSecondsSinceUnixEpoch` from the given `SystemTime`, if it is not before
    /// the unix epoch, or too large to be represented.
    pub fn from_system_time(time: SystemTime) -> Option<Self> {
        let duration = time.duration_since(UNIX_EPOCH).ok()?;
        let millis = duration.as_millis().try_into().ok()?;
        Some(Self(millis))
    }

    /// The current system time in milliseconds since the unix epoch.
    pub fn now() -> Self {
        Self::from_system_time(SystemTime::now()).unwrap()
    }

    /// Creates a new `SystemTime` from `self`, if it can be represented.
    pub fn to_system_time(self) -> Option<SystemTime> {
        UNIX_EPOCH.checked_add(Duration::from_millis(self.0.into()))
    }

    /// Get the time since the unix epoch in milliseconds.
    pub fn get(&self) -> UInt {
        self.0
    }

    /// Get time since the unix epoch in seconds.
    pub fn as_secs(&self) -> UInt {
        self.0 / uint!(1000)
    }
}

/// A timestamp represented as the number of seconds since the unix epoch.
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize)]
#[allow(clippy::exhaustive_structs)]
#[serde(transparent)]
pub struct SecondsSinceUnixEpoch(pub UInt);

impl SecondsSinceUnixEpoch {
    /// Creates a new `MilliSecondsSinceUnixEpoch` from the given `SystemTime`, if it is not before
    /// the unix epoch, or too large to be represented.
    pub fn from_system_time(time: SystemTime) -> Option<Self> {
        let duration = time.duration_since(UNIX_EPOCH).ok()?;
        let millis = duration.as_secs().try_into().ok()?;
        Some(Self(millis))
    }

    /// The current system-time as seconds since the unix epoch.
    pub fn now() -> Self {
        Self::from_system_time(SystemTime::now()).unwrap()
    }

    /// Creates a new `SystemTime` from `self`, if it can be represented.
    pub fn to_system_time(self) -> Option<SystemTime> {
        UNIX_EPOCH.checked_add(Duration::from_secs(self.0.into()))
    }

    /// Get time since the unix epoch in seconds.
    pub fn get(&self) -> UInt {
        self.0
    }
}

#[cfg(test)]
mod tests {
    use std::time::{Duration, UNIX_EPOCH};

    use js_int::uint;
    use matches::assert_matches;
    use serde::{Deserialize, Serialize};
    use serde_json::json;

    use super::{MilliSecondsSinceUnixEpoch, SecondsSinceUnixEpoch};

    #[derive(Clone, Debug, Deserialize, Serialize)]
    struct SystemTimeTest {
        millis: MilliSecondsSinceUnixEpoch,
        secs: SecondsSinceUnixEpoch,
    }

    #[test]
    fn deserialize() {
        let json = json!({ "millis": 3000, "secs": 60 });

        assert_matches!(
            serde_json::from_value::<SystemTimeTest>(json),
            Ok(SystemTimeTest { millis, secs })
            if millis.to_system_time() == Some(UNIX_EPOCH + Duration::from_millis(3000))
                && secs.to_system_time() == Some(UNIX_EPOCH + Duration::from_secs(60))
        );
    }

    #[test]
    fn serialize() {
        let request = SystemTimeTest {
            millis: MilliSecondsSinceUnixEpoch::from_system_time(UNIX_EPOCH + Duration::new(2, 0))
                .unwrap(),
            secs: SecondsSinceUnixEpoch(uint!(0)),
        };
        assert_matches!(
            serde_json::to_value(&request),
            Ok(value) if value == json!({ "millis": 2000, "secs": 0 })
        );
    }
}
<|endoftext|>