<commit_before>/// Serial port management module supporting one
/// writer and multiple readers
///
/// Clients can lock a port for writing, but
/// subscribe to data from multiple ports for reads

use std::collections::HashSet;
use std::sync::mpsc::{Receiver, TryRecvError};
use std::thread;
use std::time::Duration;

use base64;

use serial_support::errors::*;
use serial_support::messages::*;
use serial_support::sub_manager::*;
use serial_support::port_manager::*;
use serial_support::writelock_manager::*;
use serial_support::common::*;

/// Manager manages connection state.
/// TODO Needs a better name
pub struct Manager {
  /// Manage write lock status
  writelock_manager: WriteLockManager,
  /// Manage ports
  port_manager: PortManager,
  /// Manage subscriptions
  sub_manager: SubscriptionManager,
  /// Receiver for serial requests
  receiver: Receiver<(String, SerialRequest)>,
  /// Receiver for response subscription requests
  subsc_receiver: SubscReceiver,
}

impl Manager {
  ///Constructor
  pub fn new(
    receiver: Receiver<(String, SerialRequest)>,
    subsc_receiver: SubscReceiver,
  ) -> Manager {
    Manager {
      writelock_manager: WriteLockManager::new(),
      port_manager: PortManager::new(),
      sub_manager: SubscriptionManager::new(),
      receiver: receiver,
      subsc_receiver: subsc_receiver,
    }
  }

  ///Spawn an instance in a new thread.
  pub fn spawn(
    receiver: Receiver<(String, SerialRequest)>,
    subsc_receiver: SubscReceiver,
  ) -> thread::JoinHandle<()> {
    thread::spawn(move || { Manager::new(receiver, subsc_receiver).run(); })
  }

  /// Main loop
  fn run(&mut self) {

    // Bad ports we couldn't read from
    // A set of SerialResponse::Errors built from
    // from the serial port read/write error responses
    let mut bad_ports = HashSet::<String>::new();
    // Shutdown flag
    let mut shutdown = false;
    // Check about 30 times a second
    let sleep_dur = Duration::from_millis(33);

    while !shutdown {
      // Sleep for a little bit to avoid pegging cpu
      // TODO Dynamic timing loop
      thread::sleep(sleep_dur);

      // Handle serial operation requests
      match self.receiver.try_recv() {
        Err(e) => {
          match e {
            TryRecvError::Empty => {
              // nothing to do
            }
            TryRecvError::Disconnected => {
              // Remote end hung up, time to shutdown
              info!("Shutting down SerialPortManager");
              shutdown = true;
            }
          }
        }
        Ok(req) => self.handle_serial_request(&req.0, req.1),
      }

      // Check for new data on each port
      for (port_name, result) in self.port_manager.read_all_ports() {
        match result {
          Ok(data) => {
            let response = match String::from_utf8(data) {
              // We need to send as binary
              Err(e) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: base64::encode(&e.into_bytes()),
                  base64: Some(true),
                }
              }
              Ok(s) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: s,
                  base64: Some(false),
                }
              }
            };
            self.broadcast_message_for_port(&port_name, response);
          }
          // Send data reads
          Err(e) => {
            warn!("Error reading port!");
            warn!("{}", e);
            bad_ports.insert(port_name);
          }
        }
      }

      //Handle write requests
      let mut recv_count = 0;
      let mut do_recv = true;
      while do_recv {
        recv_count += 1;
        match self.subsc_receiver.try_recv() {
          Ok(sub_request) => self.sub_manager.add_subscription(sub_request),
          Err(e) => {
            match e {
              TryRecvError::Disconnected => {
                // Does this mean all senders have disconnected?
                // Or just one?
                debug!("Got disconnected when trying to get serial request");
              }
              TryRecvError::Empty => do_recv = false,
            }
          }
        }

        // Don't let this loop get stuck here forever
        // starving the outer loop
        if recv_count == 50 {
          do_recv = false;
        }
      }

      // Cleanup bad serial ports that failed read or write
      // We remove them from everything before
      self.cleanup_bad_ports(&bad_ports);
      bad_ports.clear();
    }
  }

  /// Handles and dispatches SerialRequest sent by
  /// the channel
  fn handle_serial_request(&mut self, sub_id: &String, msg: SerialRequest) {
    let response = match msg {
      SerialRequest::Open { port } => self.handle_open_port(sub_id, port),
      SerialRequest::WriteLock { port } => self.handle_write_lock(sub_id, port),
      SerialRequest::ReleaseWriteLock { port } => self.handle_release_write_lock(sub_id, port),
      SerialRequest::Write { port, data, base64 } => {
        self.handle_write_port(sub_id, port, data, base64.unwrap_or(false))
      }
      SerialRequest::Close { port } => self.handle_close_port(sub_id, port),
      SerialRequest::List {} => self.handle_list_ports(sub_id),
    };
    match response {
      Ok(response) => {
        // Send acknowledge?
      }
      Err(e) => {
        warn!("Error '{}' occured handling serial request message", e);
        // Send error?
        self.send_message(&sub_id, to_serial_response_error(e));
      }
    }
  }

  /// Handle write port requests
  fn handle_write_port(
    &mut self,
    sub_id: &String,
    port_name: String,
    data: String,
    base_64: bool,
  ) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self.check_owns_writelock(&port_name, &sub_id)?;
    match base_64 {
      true => {
        base64::decode(&data)
          .map_err(|e| ErrorKind::Base64(e).into())
          .and_then(|d| self.port_manager.write_port(&port_name, &d))
          .map(|_| self.send_message(&sub_id, SerialResponse::Wrote { port: port_name }),)
      }
      false => {
        self
          .port_manager
          .write_port(&port_name, data.as_bytes())
          .map(|_| self.send_message(&sub_id, SerialResponse::Wrote { port: port_name }),)
      }
    }
  }

  /// Handle write lock requests
  fn handle_write_lock(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self
      .writelock_manager
      .lock_port(&port_name, &sub_id)
      .map(|_| self.send_message(&sub_id, SerialResponse::WriteLocked { port: port_name }),)
  }

  /// Handle write requests
  fn handle_release_write_lock(
    &mut self,
    sub_id: &String,
    port_name: Option<String>,
  ) -> Result<()> {
    self.check_sub_id(sub_id)?;
    match port_name {
      None => {
        self.writelock_manager.unlock_all_ports_for_sub(&sub_id);
        Ok(
          self.send_message(
            &sub_id,
            SerialResponse::WriteLockReleased { port: port_name },
          ),
        )
      }
      Some(port_name) => {
        self
          .writelock_manager
          .unlock_port(&port_name, &sub_id)
          .map(
            |_| {
              self.send_message(
                &sub_id,
                SerialResponse::WriteLockReleased { port: Some(port_name) },
              )
            },
          )
      }
    }
  }

  /// Handle open port requests
  fn handle_open_port(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self.port_manager.open_port(&port_name)?;
    self
      .sub_manager
      .add_port(&sub_id, &port_name)
      .map(|_| self.send_message(&sub_id, SerialResponse::Opened { port: port_name }),)
  }

  /// Handle list ports request
  fn handle_list_ports(&mut self, sub_id: &String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    let port_names: Result<Vec<String>> =
      self
        .port_manager
        .list_ports()
        .map(|v| v.iter().map(|v| v.port_name.clone()).collect());
    let resp = port_names.map(|pns| SerialResponse::List { ports: pns });
    self.send_message(
      &sub_id,
      resp.unwrap_or(
        SerialResponse::Error {
          display: "".to_string(),
          description: "".to_string(),
        },
      ),
    );
    Ok(())
  }

  /// Handle close port requests
  fn handle_close_port(&mut self, sub_id: &String, port_name: Option<String>) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    // Unsubscribe from ports
    match port_name {
      None => {
        self.sub_manager.clear_ports(Some(&sub_id));
        Ok(())
        },
      Some(ref pn) => self.sub_manager.remove_port(&sub_id, &pn)
    }?;
    // Remove write locks
    match port_name {
      None => {
        self.writelock_manager.unlock_all_ports_for_sub(sub_id);
        Ok(())
      },
      Some(ref pn) => self.writelock_manager.unlock_port(&pn, sub_id)
    }?;
    // TODO: Ensure close messages are sent, etc

    // Close ports with no subscribers
    let open_ports = self.port_manager.open_ports();
    let subscribed_ports = self.sub_manager.subscribed_ports();
    let ports_with_no_subs = open_ports.difference(&subscribed_ports);

    // For each open port that isn't subscribed,
    for port_to_close in ports_with_no_subs {
      // close it
      self.port_manager.close_port(&port_to_close);
      // remove the write lock
      self.writelock_manager.clear_lock(&port_to_close);
      // Let them know its closed
      let close_resp = SerialResponse::Closed { port: port_to_close.clone() };
      self.send_message(&sub_id, close_resp);
    }
    Ok(())
  }

  /// Cleanup any bad ports
  /// Clears out the passed in hash set
  fn cleanup_bad_ports(&mut self, bad_ports: &HashSet<String>) {
    for port_name in bad_ports.iter() {
      // Tell everyone port is sick
      let err_resp = SerialResponse::Error {
        display: format!("Error reading from port '{}', closing!", port_name),
        description: format!(
          "An error occured while trying to read data from '{}'",
          port_name
        ),
      };
      self.broadcast_message_for_port(port_name, err_resp);
      // Tesll everyone the sick ports were closed
      let close_resp = SerialResponse::Closed { port: port_name.clone() };
      self.broadcast_message_for_port(port_name, close_resp);
      // Close bad ports
      self.port_manager.close_port(port_name);
      // Remove write locks on bad ports
      self.writelock_manager.clear_lock(port_name);
      // Remove bad ports from subscriptions
      self.sub_manager.remove_port_from_all(port_name);
    }
    // Clear set of bad_ports
  }

  pub fn send_message(&mut self, sub_id: &String, msg: SerialResponse) {
    match self.sub_manager.send_message(sub_id, msg) {
      Err(e) => {
        warn!("Error sending serial response to sub_id '{}'", sub_id);
        let mut bad_subs = Vec::new();
        bad_subs.push(e);
        self.cleanup_bad_subs(bad_subs);
      }
      _ => {}
    }
  }

  /// Broadcast a message to all subscribers and
  /// then cleanup any subs that errored
  pub fn broadcast_message(&mut self, msg: SerialResponse) {
    let bad_subs = self.sub_manager.broadcast_message(msg);
    self.cleanup_bad_subs(bad_subs);
  }

  /// Broadcast a message to all subscribers interested in the
  /// given port and then cleanup any subs that errored
  pub fn broadcast_message_for_port(&mut self, port_name: &String, msg: SerialResponse) {
    let bad_subs = self
      .sub_manager
      .broadcast_message_for_port(port_name, msg);
    self.cleanup_bad_subs(bad_subs);
  }

  /// Cleanup any bad subs where a message send failed
  fn cleanup_bad_subs(&mut self, bad_subs: Vec<Error>) {
    for e in bad_subs {
      match e {
        Error(ErrorKind::SubscriberSendError(sub_id), _) => {
          self.sub_manager.end_subscription(&sub_id);
          self.writelock_manager.unlock_all_ports_for_sub(&sub_id);
        }
        _ => { /*nop*/ }
      }
    }
  }

  fn check_sub_id(&self, sub_id: &String) -> Result<()> {
    self.sub_manager.check_subscription_exists(sub_id)
  }

  fn check_owns_writelock(&self, port_name: &String, sub_id: &String) -> Result<()> {
    self
      .writelock_manager
      .check_owns_write_lock(port_name, sub_id)
  }
}
<commit_msg>Change order of operations, fix response error<commit_after>/// Serial port management module supporting one
/// writer and multiple readers
///
/// Clients can lock a port for writing, but
/// subscribe to data from multiple ports for reads

use std::collections::HashSet;
use std::sync::mpsc::{Receiver, TryRecvError};
use std::thread;
use std::time::Duration;

use base64;

use serial_support::errors::*;
use serial_support::messages::*;
use serial_support::sub_manager::*;
use serial_support::port_manager::*;
use serial_support::writelock_manager::*;
use serial_support::common::*;

/// Manager manages connection state.
/// TODO Needs a better name
pub struct Manager {
  /// Manage write lock status
  writelock_manager: WriteLockManager,
  /// Manage ports
  port_manager: PortManager,
  /// Manage subscriptions
  sub_manager: SubscriptionManager,
  /// Receiver for serial requests
  receiver: Receiver<(String, SerialRequest)>,
  /// Receiver for response subscription requests
  subsc_receiver: SubscReceiver,
}

impl Manager {
  ///Constructor
  pub fn new(
    receiver: Receiver<(String, SerialRequest)>,
    subsc_receiver: SubscReceiver,
  ) -> Manager {
    Manager {
      writelock_manager: WriteLockManager::new(),
      port_manager: PortManager::new(),
      sub_manager: SubscriptionManager::new(),
      receiver: receiver,
      subsc_receiver: subsc_receiver,
    }
  }

  ///Spawn an instance in a new thread.
  pub fn spawn(
    receiver: Receiver<(String, SerialRequest)>,
    subsc_receiver: SubscReceiver,
  ) -> thread::JoinHandle<()> {
    thread::spawn(move || { Manager::new(receiver, subsc_receiver).run(); })
  }

  /// Main loop
  fn run(&mut self) {

    // Bad ports we couldn't read from
    // A set of SerialResponse::Errors built from
    // from the serial port read/write error responses
    let mut bad_ports = HashSet::<String>::new();
    // Shutdown flag
    let mut shutdown = false;
    // Check about 30 times a second
    let sleep_dur = Duration::from_millis(33);

    while !shutdown {
      // Sleep for a little bit to avoid pegging cpu
      // TODO Dynamic timing loop
      thread::sleep(sleep_dur);

      // Handle serial operation requests
      match self.receiver.try_recv() {
        Err(e) => {
          match e {
            TryRecvError::Empty => {
              // nothing to do
            }
            TryRecvError::Disconnected => {
              // Remote end hung up, time to shutdown
              info!("Shutting down SerialPortManager");
              shutdown = true;
            }
          }
        }
        Ok(req) => self.handle_serial_request(&req.0, req.1),
      }

      // Check for new data on each port
      for (port_name, result) in self.port_manager.read_all_ports() {
        match result {
          Ok(data) => {
            let response = match String::from_utf8(data) {
              // We need to send as binary
              Err(e) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: base64::encode(&e.into_bytes()),
                  base64: Some(true),
                }
              }
              Ok(s) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: s,
                  base64: Some(false),
                }
              }
            };
            self.broadcast_message_for_port(&port_name, response);
          }
          // Send data reads
          Err(e) => {
            warn!("Error reading port!");
            warn!("{}", e);
            bad_ports.insert(port_name);
          }
        }
      }

      //Handle write requests
      let mut recv_count = 0;
      let mut do_recv = true;
      while do_recv {
        recv_count += 1;
        match self.subsc_receiver.try_recv() {
          Ok(sub_request) => self.sub_manager.add_subscription(sub_request),
          Err(e) => {
            match e {
              TryRecvError::Disconnected => {
                // Does this mean all senders have disconnected?
                // Or just one?
                debug!("Got disconnected when trying to get serial request");
              }
              TryRecvError::Empty => do_recv = false,
            }
          }
        }

        // Don't let this loop get stuck here forever
        // starving the outer loop
        if recv_count == 50 {
          do_recv = false;
        }
      }

      // Cleanup bad serial ports that failed read or write
      // We remove them from everything before
      self.cleanup_bad_ports(&bad_ports);
      bad_ports.clear();
    }
  }

  /// Handles and dispatches SerialRequest sent by
  /// the channel
  fn handle_serial_request(&mut self, sub_id: &String, msg: SerialRequest) {
    let response = match msg {
      SerialRequest::Open { port } => self.handle_open_port(sub_id, port),
      SerialRequest::WriteLock { port } => self.handle_write_lock(sub_id, port),
      SerialRequest::ReleaseWriteLock { port } => self.handle_release_write_lock(sub_id, port),
      SerialRequest::Write { port, data, base64 } => {
        self.handle_write_port(sub_id, port, data, base64.unwrap_or(false))
      }
      SerialRequest::Close { port } => self.handle_close_port(sub_id, port),
      SerialRequest::List {} => self.handle_list_ports(sub_id),
    };
    match response {
      Ok(response) => {
        // Send acknowledge?
      }
      Err(e) => {
        warn!("Error '{}' occured handling serial request message", e);
        // Send error?
        self.send_message(&sub_id, to_serial_response_error(e));
      }
    }
  }

  /// Handle write port requests
  fn handle_write_port(
    &mut self,
    sub_id: &String,
    port_name: String,
    data: String,
    base_64: bool,
  ) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self.check_owns_writelock(&port_name, &sub_id)?;
    match base_64 {
      true => {
        base64::decode(&data)
          .map_err(|e| ErrorKind::Base64(e).into())
          .and_then(|d| self.port_manager.write_port(&port_name, &d))
          .map(|_| self.send_message(&sub_id, SerialResponse::Wrote { port: port_name }),)
      }
      false => {
        self
          .port_manager
          .write_port(&port_name, data.as_bytes())
          .map(|_| self.send_message(&sub_id, SerialResponse::Wrote { port: port_name }),)
      }
    }
  }

  /// Handle write lock requests
  fn handle_write_lock(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self
      .writelock_manager
      .lock_port(&port_name, &sub_id)
      .map(|_| self.send_message(&sub_id, SerialResponse::WriteLocked { port: port_name }),)
  }

  /// Handle write requests
  fn handle_release_write_lock(
    &mut self,
    sub_id: &String,
    port_name: Option<String>,
  ) -> Result<()> {
    self.check_sub_id(sub_id)?;
    match port_name {
      None => {
        self.writelock_manager.unlock_all_ports_for_sub(&sub_id);
        Ok(
          self.send_message(
            &sub_id,
            SerialResponse::WriteLockReleased { port: port_name },
          ),
        )
      }
      Some(port_name) => {
        self
          .writelock_manager
          .unlock_port(&port_name, &sub_id)
          .map(
            |_| {
              self.send_message(
                &sub_id,
                SerialResponse::WriteLockReleased { port: Some(port_name) },
              )
            },
          )
      }
    }
  }

  /// Handle open port requests
  fn handle_open_port(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self.port_manager.open_port(&port_name)?;
    self
      .sub_manager
      .add_port(&sub_id, &port_name)
      .map(|_| self.send_message(&sub_id, SerialResponse::Opened { port: port_name }),)
  }

  /// Handle list ports request
  fn handle_list_ports(&mut self, sub_id: &String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    let port_names: Result<Vec<String>> =
      self
        .port_manager
        .list_ports()
        .map(|v| v.iter().map(|v| v.port_name.clone()).collect());
    let resp = port_names.map(|pns| SerialResponse::List { ports: pns });
    self.send_message(
      &sub_id,
      resp.unwrap_or(
        SerialResponse::Error {
          display: "".to_string(),
          description: "".to_string(),
        },
      ),
    );
    Ok(())
  }

  /// Handle close port requests
  fn handle_close_port(&mut self, sub_id: &String, port_name: Option<String>) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    // Unsubscribe from ports
    match port_name {
      None => {
        self.sub_manager.clear_ports(Some(&sub_id));
        Ok(())
        },
      Some(ref pn) => self.sub_manager.remove_port(&sub_id, &pn)
    }?;
    // Remove write locks
    match port_name {
      None => {
        self.writelock_manager.unlock_all_ports_for_sub(sub_id);
        Ok(())
      },
      Some(ref pn) => self.writelock_manager.unlock_port(&pn, sub_id)
    }?;
    // TODO: Ensure close messages are sent, etc

    // Close ports with no subscribers
    let open_ports = self.port_manager.open_ports();
    let subscribed_ports = self.sub_manager.subscribed_ports();
    let ports_with_no_subs = open_ports.difference(&subscribed_ports);

    // For each open port that isn't subscribed,
    for port_to_close in ports_with_no_subs {
      // close it
      self.port_manager.close_port(&port_to_close);
      // remove the write lock
      self.writelock_manager.clear_lock(&port_to_close);
      // Let them know its closed
      let close_resp = SerialResponse::Closed { port: port_to_close.clone() };
      self.send_message(&sub_id, close_resp);
    }
    Ok(())
  }

  /// Cleanup any bad ports
  /// Clears out the passed in hash set
  fn cleanup_bad_ports(&mut self, bad_ports: &HashSet<String>) {
    for port_name in bad_ports.iter() {
      // Close bad ports
      self.port_manager.close_port(port_name);
      // Remove write locks on bad ports
      self.writelock_manager.clear_lock(port_name);
      // Remove bad ports from subscriptions
      self.sub_manager.remove_port_from_all(port_name);
      // Tell everyone port is sick
      let err_resp = to_serial_response_error(ErrorKind::PortReadError(port_name.to_owned()).into());
      self.broadcast_message_for_port(port_name, err_resp);
      // Tesll everyone the sick ports were closed
      let close_resp = SerialResponse::Closed { port: port_name.clone() };
      self.broadcast_message_for_port(port_name, close_resp);
    }
    // Clear set of bad_ports
  }

  pub fn send_message(&mut self, sub_id: &String, msg: SerialResponse) {
    match self.sub_manager.send_message(sub_id, msg) {
      Err(e) => {
        warn!("Error sending serial response to sub_id '{}'", sub_id);
        let mut bad_subs = Vec::new();
        bad_subs.push(e);
        self.cleanup_bad_subs(bad_subs);
      }
      _ => {}
    }
  }

  /// Broadcast a message to all subscribers and
  /// then cleanup any subs that errored
  pub fn broadcast_message(&mut self, msg: SerialResponse) {
    let bad_subs = self.sub_manager.broadcast_message(msg);
    self.cleanup_bad_subs(bad_subs);
  }

  /// Broadcast a message to all subscribers interested in the
  /// given port and then cleanup any subs that errored
  pub fn broadcast_message_for_port(&mut self, port_name: &String, msg: SerialResponse) {
    let bad_subs = self
      .sub_manager
      .broadcast_message_for_port(port_name, msg);
    self.cleanup_bad_subs(bad_subs);
  }

  /// Cleanup any bad subs where a message send failed
  fn cleanup_bad_subs(&mut self, bad_subs: Vec<Error>) {
    for e in bad_subs {
      match e {
        Error(ErrorKind::SubscriberSendError(sub_id), _) => {
          self.sub_manager.end_subscription(&sub_id);
          self.writelock_manager.unlock_all_ports_for_sub(&sub_id);
        }
        _ => { /*nop*/ }
      }
    }
  }

  fn check_sub_id(&self, sub_id: &String) -> Result<()> {
    self.sub_manager.check_subscription_exists(sub_id)
  }

  fn check_owns_writelock(&self, port_name: &String, sub_id: &String) -> Result<()> {
    self
      .writelock_manager
      .check_owns_write_lock(port_name, sub_id)
  }
}
<|endoftext|>


<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) 2015 Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// * (c) 2020 nicoo            <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use num_traits::{
    cast::{AsPrimitive, ToPrimitive},
    identities::{One, Zero},
    int::PrimInt,
    ops::wrapping::{WrappingMul, WrappingNeg, WrappingSub},
};
use std::fmt::{Debug, Display};
use std::mem::swap;

// This is incorrectly reported as dead code,
//  presumably when included in build.rs.
#[allow(dead_code)]
pub(crate) fn gcd(mut a: u64, mut b: u64) -> u64 {
    while b > 0 {
        a %= b;
        swap(&mut a, &mut b);
    }
    a
}

pub(crate) trait Arithmetic: Copy + Sized {
    // The type of integers mod m, in some opaque representation
    type ModInt: Copy + Sized + Eq;

    fn new(m: u64) -> Self;
    fn modulus(&self) -> u64;
    fn from_u64(&self, n: u64) -> Self::ModInt;
    fn to_u64(&self, n: Self::ModInt) -> u64;
    fn add(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt;
    fn mul(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt;

    fn pow(&self, mut a: Self::ModInt, mut b: u64) -> Self::ModInt {
        let (_a, _b) = (a, b);
        let mut result = self.one();
        while b > 0 {
            if b & 1 != 0 {
                result = self.mul(result, a);
            }
            a = self.mul(a, a);
            b >>= 1;
        }

        // Check that r (reduced back to the usual representation) equals
        //  a^b % n, unless the latter computation overflows
        // Temporarily commented-out, as there u64::checked_pow is not available
        //  on the minimum supported Rust version, nor is an appropriate method
        //  for compiling the check conditionally.
        //debug_assert!(self
        //    .to_u64(_a)
        //    .checked_pow(_b as u32)
        //    .map(|r| r % self.modulus() == self.to_u64(result))
        //    .unwrap_or(true));

        result
    }

    fn one(&self) -> Self::ModInt {
        self.from_u64(1)
    }
    fn minus_one(&self) -> Self::ModInt {
        self.from_u64(self.modulus() - 1)
    }
    fn zero(&self) -> Self::ModInt {
        self.from_u64(0)
    }
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct Montgomery<T: DoubleInt> {
    a: T,
    n: T,
}

impl<T: DoubleInt> Montgomery<T> {
    /// computes x/R mod n efficiently
    fn reduce(&self, x: T::Double) -> T {
        let t_bits = T::zero().count_zeros() as usize;

        debug_assert!(x < (self.n.as_double()) << t_bits);
        // TODO: optimiiiiiiise
        let Montgomery { a, n } = self;
        let m = T::from_double(x).wrapping_mul(a);
        let nm = (n.as_double()) * (m.as_double());
        let (xnm, overflow) = x.overflowing_add_(nm); // x + n*m
        debug_assert_eq!(
            xnm % (T::Double::one() << T::zero().count_zeros() as usize),
            T::Double::zero()
        );

        // (x + n*m) / R
        // in case of overflow, this is (2¹²⁸ + xnm)/2⁶⁴ - n = xnm/2⁶⁴ + (2⁶⁴ - n)
        let y = T::from_double(xnm >> t_bits)
            + if !overflow {
                T::zero()
            } else {
                n.wrapping_neg()
            };

        if y >= *n {
            y - *n
        } else {
            y
        }
    }
}

impl<T: DoubleInt> Arithmetic for Montgomery<T> {
    // Montgomery transform, R=2⁶⁴
    // Provides fast arithmetic mod n (n odd, u64)
    type ModInt = T;

    fn new(n: u64) -> Self {
        debug_assert!(T::zero().count_zeros() >= 64 || n < (1 << T::zero().count_zeros() as usize));
        let n = T::from_u64(n);
        let a = modular_inverse(n).wrapping_neg();
        debug_assert_eq!(n.wrapping_mul(&a), T::one().wrapping_neg());
        Montgomery { a, n }
    }

    fn modulus(&self) -> u64 {
        self.n.as_()
    }

    fn from_u64(&self, x: u64) -> Self::ModInt {
        // TODO: optimise!
        debug_assert!(x < self.n.as_());
        let r = T::from_double(
            ((T::Double::from(x)) << T::zero().count_zeros() as usize) % self.n.as_double(),
        );
        debug_assert_eq!(x, self.to_u64(r));
        r
    }

    fn to_u64(&self, n: Self::ModInt) -> u64 {
        self.reduce(n.as_double()).as_()
    }

    fn add(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt {
        let (r, overflow) = a.overflowing_add_(b);

        // In case of overflow, a+b = 2⁶⁴ + r = (2⁶⁴ - n) + r (working mod n)
        let r = if !overflow {
            r
        } else {
            r + self.n.wrapping_neg()
        };

        // Normalise to [0; n[
        let r = if r < self.n { r } else { r - self.n };

        // Check that r (reduced back to the usual representation) equals
        // a+b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2 =
                (((a_r as u128) + (b_r as u128)) % <T as AsPrimitive<u128>>::as_(self.n)) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} + {} = [{}] + [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }

    fn mul(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt {
        let r = self.reduce(a.as_double() * b.as_double());

        // Check that r (reduced back to the usual representation) equals
        // a*b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2: u64 = ((T::Double::from(a_r) * T::Double::from(b_r)) % self.n.as_double())
                .to_u64()
                .unwrap();
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} * {} = [{}] * [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }
}

// NOTE: Trait can be removed once num-traits adds a similar one;
//       see https://github.com/rust-num/num-traits/issues/168
pub(crate) trait OverflowingAdd: Sized {
    fn overflowing_add_(self, n: Self) -> (Self, bool);
}
impl OverflowingAdd for u32 {
    fn overflowing_add_(self, n: Self) -> (Self, bool) {
        self.overflowing_add(n)
    }
}
impl OverflowingAdd for u64 {
    fn overflowing_add_(self, n: Self) -> (Self, bool) {
        self.overflowing_add(n)
    }
}
impl OverflowingAdd for u128 {
    fn overflowing_add_(self, n: Self) -> (Self, bool) {
        self.overflowing_add(n)
    }
}

pub(crate) trait Int:
    AsPrimitive<u64>
    + AsPrimitive<u128>
    + Display
    + Debug
    + PrimInt
    + OverflowingAdd
    + WrappingNeg
    + WrappingSub
    + WrappingMul
{
}

pub(crate) trait DoubleInt: Int {
    type Double: Int + From<u64> + ToPrimitive;

    fn as_double(self) -> Self::Double;
    fn from_double(n: Self::Double) -> Self;
    fn from_u64(n: u64) -> Self;
}

impl Int for u32 {}
impl Int for u64 {}
impl Int for u128 {}

impl DoubleInt for u64 {
    type Double = u128;

    fn as_double(self) -> u128 {
        self as _
    }
    fn from_double(n: u128) -> u64 {
        n as _
    }
    fn from_u64(n: u64) -> u64 {
        n
    }
}
impl DoubleInt for u32 {
    type Double = u64;

    fn as_double(self) -> u64 {
        self as _
    }
    fn from_double(n: u64) -> u32 {
        n as _
    }
    fn from_u64(n: u64) -> u32 {
        n as _
    }
}

// extended Euclid algorithm
// precondition: a is odd
pub(crate) fn modular_inverse<T: Int>(a: T) -> T {
    let zero = T::zero();
    let one = T::one();
    debug_assert!(a % (one + one) == one, "{:?} is not odd", a);

    let mut t = zero;
    let mut newt = one;
    let mut r = zero;
    let mut newr = a;

    while newr != zero {
        let quot = if r == zero {
            // special case when we're just starting out
            // This works because we know that
            // a does not divide 2^64, so floor(2^64 / a) == floor((2^64-1) / a);
            T::max_value()
        } else {
            r
        } / newr;

        let newtp = t.wrapping_sub(&quot.wrapping_mul(&newt));
        t = newt;
        newt = newtp;

        let newrp = r.wrapping_sub(&quot.wrapping_mul(&newr));
        r = newr;
        newr = newrp;
    }

    debug_assert_eq!(r, one);
    t
}

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

    fn test_inverter<T: Int>() {
        // All odd integers from 1 to 20 000
        let one = T::from(1).unwrap();
        let two = T::from(2).unwrap();
        let mut test_values = (0..10_000)
            .map(|i| T::from(i).unwrap())
            .map(|i| two * i + one);

        assert!(test_values.all(|x| x.wrapping_mul(&modular_inverse(x)) == one));
    }

    #[test]
    fn test_inverter_u32() {
        test_inverter::<u32>()
    }

    #[test]
    fn test_inverter_u64() {
        test_inverter::<u64>()
    }

    fn test_add<A: Arithmetic>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = A::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    println!("{n:?}, {x:?}, {y:?}", n = n, x = x, y = y);
                    assert_eq!((x + y) % n, m.to_u64(m.add(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_add_m32() {
        test_add::<Montgomery<u32>>()
    }

    #[test]
    fn test_add_m64() {
        test_add::<Montgomery<u64>>()
    }

    fn test_mult<A: Arithmetic>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = A::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    assert_eq!((x * y) % n, m.to_u64(m.mul(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_mult_m32() {
        test_mult::<Montgomery<u32>>()
    }

    #[test]
    fn test_mult_m64() {
        test_mult::<Montgomery<u64>>()
    }

    fn test_roundtrip<A: Arithmetic>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = A::new(n);
            for x in 0..n {
                let x_ = m.from_u64(x);
                assert_eq!(x, m.to_u64(x_));
            }
        }
    }

    #[test]
    fn test_roundtrip_m32() {
        test_roundtrip::<Montgomery<u32>>()
    }

    #[test]
    fn test_roundtrip_m64() {
        test_roundtrip::<Montgomery<u64>>()
    }
}
<commit_msg>factor::numeric: Refactor away the use of {To,From}Primitives<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) 2015 Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// * (c) 2020 nicoo            <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use num_traits::{
    identities::{One, Zero},
    int::PrimInt,
    ops::wrapping::{WrappingMul, WrappingNeg, WrappingSub},
};
use std::fmt::{Debug, Display};
use std::mem::swap;

// This is incorrectly reported as dead code,
//  presumably when included in build.rs.
#[allow(dead_code)]
pub(crate) fn gcd(mut a: u64, mut b: u64) -> u64 {
    while b > 0 {
        a %= b;
        swap(&mut a, &mut b);
    }
    a
}

pub(crate) trait Arithmetic: Copy + Sized {
    // The type of integers mod m, in some opaque representation
    type ModInt: Copy + Sized + Eq;

    fn new(m: u64) -> Self;
    fn modulus(&self) -> u64;
    fn from_u64(&self, n: u64) -> Self::ModInt;
    fn to_u64(&self, n: Self::ModInt) -> u64;
    fn add(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt;
    fn mul(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt;

    fn pow(&self, mut a: Self::ModInt, mut b: u64) -> Self::ModInt {
        let (_a, _b) = (a, b);
        let mut result = self.one();
        while b > 0 {
            if b & 1 != 0 {
                result = self.mul(result, a);
            }
            a = self.mul(a, a);
            b >>= 1;
        }

        // Check that r (reduced back to the usual representation) equals
        //  a^b % n, unless the latter computation overflows
        // Temporarily commented-out, as there u64::checked_pow is not available
        //  on the minimum supported Rust version, nor is an appropriate method
        //  for compiling the check conditionally.
        //debug_assert!(self
        //    .to_u64(_a)
        //    .checked_pow(_b as u32)
        //    .map(|r| r % self.modulus() == self.to_u64(result))
        //    .unwrap_or(true));

        result
    }

    fn one(&self) -> Self::ModInt {
        self.from_u64(1)
    }
    fn minus_one(&self) -> Self::ModInt {
        self.from_u64(self.modulus() - 1)
    }
    fn zero(&self) -> Self::ModInt {
        self.from_u64(0)
    }
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct Montgomery<T: DoubleInt> {
    a: T,
    n: T,
}

impl<T: DoubleInt> Montgomery<T> {
    /// computes x/R mod n efficiently
    fn reduce(&self, x: T::Double) -> T {
        let t_bits = T::zero().count_zeros() as usize;

        debug_assert!(x < (self.n.as_double()) << t_bits);
        // TODO: optimiiiiiiise
        let Montgomery { a, n } = self;
        let m = T::from_double(x).wrapping_mul(a);
        let nm = (n.as_double()) * (m.as_double());
        let (xnm, overflow) = x.overflowing_add_(nm); // x + n*m
        debug_assert_eq!(
            xnm % (T::Double::one() << T::zero().count_zeros() as usize),
            T::Double::zero()
        );

        // (x + n*m) / R
        // in case of overflow, this is (2¹²⁸ + xnm)/2⁶⁴ - n = xnm/2⁶⁴ + (2⁶⁴ - n)
        let y = T::from_double(xnm >> t_bits)
            + if !overflow {
                T::zero()
            } else {
                n.wrapping_neg()
            };

        if y >= *n {
            y - *n
        } else {
            y
        }
    }
}

impl<T: DoubleInt> Arithmetic for Montgomery<T> {
    // Montgomery transform, R=2⁶⁴
    // Provides fast arithmetic mod n (n odd, u64)
    type ModInt = T;

    fn new(n: u64) -> Self {
        debug_assert!(T::zero().count_zeros() >= 64 || n < (1 << T::zero().count_zeros() as usize));
        let n = T::from_u64(n);
        let a = modular_inverse(n).wrapping_neg();
        debug_assert_eq!(n.wrapping_mul(&a), T::one().wrapping_neg());
        Montgomery { a, n }
    }

    fn modulus(&self) -> u64 {
        self.n.as_u64()
    }

    fn from_u64(&self, x: u64) -> Self::ModInt {
        // TODO: optimise!
        debug_assert!(x < self.n.as_u64());
        let r = T::from_double(
            ((T::Double::from_u64(x)) << T::zero().count_zeros() as usize) % self.n.as_double(),
        );
        debug_assert_eq!(x, self.to_u64(r));
        r
    }

    fn to_u64(&self, n: Self::ModInt) -> u64 {
        self.reduce(n.as_double()).as_u64()
    }

    fn add(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt {
        let (r, overflow) = a.overflowing_add_(b);

        // In case of overflow, a+b = 2⁶⁴ + r = (2⁶⁴ - n) + r (working mod n)
        let r = if !overflow {
            r
        } else {
            r + self.n.wrapping_neg()
        };

        // Normalise to [0; n[
        let r = if r < self.n { r } else { r - self.n };

        // Check that r (reduced back to the usual representation) equals
        // a+b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a) as u128;
            let b_r = self.to_u64(b) as u128;
            let r_r = self.to_u64(r);
            let r_2 = ((a_r + b_r) % self.n.as_u128()) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} + {} = [{}] + [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }

    fn mul(&self, a: Self::ModInt, b: Self::ModInt) -> Self::ModInt {
        let r = self.reduce(a.as_double() * b.as_double());

        // Check that r (reduced back to the usual representation) equals
        // a*b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a) as u128;
            let b_r = self.to_u64(b) as u128;
            let r_r = self.to_u64(r);
            let r_2: u64 = ((a_r * b_r) % self.n.as_u128()) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} * {} = [{}] * [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }
}

// NOTE: Trait can be removed once num-traits adds a similar one;
//       see https://github.com/rust-num/num-traits/issues/168
pub(crate) trait OverflowingAdd: Sized {
    fn overflowing_add_(self, n: Self) -> (Self, bool);
}
impl OverflowingAdd for u32 {
    fn overflowing_add_(self, n: Self) -> (Self, bool) {
        self.overflowing_add(n)
    }
}
impl OverflowingAdd for u64 {
    fn overflowing_add_(self, n: Self) -> (Self, bool) {
        self.overflowing_add(n)
    }
}
impl OverflowingAdd for u128 {
    fn overflowing_add_(self, n: Self) -> (Self, bool) {
        self.overflowing_add(n)
    }
}

pub(crate) trait Int:
    Display + Debug + PrimInt + OverflowingAdd + WrappingNeg + WrappingSub + WrappingMul
{
    fn as_u64(&self) -> u64;
    fn from_u64(n: u64) -> Self;
    #[cfg(debug_assertions)]
    fn as_u128(&self) -> u128;
    #[cfg(debug_assertions)]
    fn from_u128(n: u64) -> Self;
}

pub(crate) trait DoubleInt: Int {
    type Double: Int;

    fn as_double(self) -> Self::Double;
    fn from_double(n: Self::Double) -> Self;
}

impl Int for u32 {
    fn as_u64(&self) -> u64 {
        *self as _
    }
    fn from_u64(n: u64) -> Self {
        n as _
    }
    #[cfg(debug_assertions)]
    fn as_u128(&self) -> u128 {
        *self as _
    }
    #[cfg(debug_assertions)]
    fn from_u128(n: u64) -> Self {
        n as _
    }
}
impl Int for u64 {
    fn as_u64(&self) -> u64 {
        *self as _
    }
    fn from_u64(n: u64) -> Self {
        n as _
    }
    #[cfg(debug_assertions)]
    fn as_u128(&self) -> u128 {
        *self as _
    }
    #[cfg(debug_assertions)]
    fn from_u128(n: u64) -> Self {
        n as _
    }
}
impl Int for u128 {
    fn as_u64(&self) -> u64 {
        *self as _
    }
    fn from_u64(n: u64) -> Self {
        n as _
    }
    #[cfg(debug_assertions)]
    fn as_u128(&self) -> u128 {
        *self as _
    }
    #[cfg(debug_assertions)]
    fn from_u128(n: u64) -> Self {
        n as _
    }
}

impl DoubleInt for u64 {
    type Double = u128;

    fn as_double(self) -> u128 {
        self as _
    }
    fn from_double(n: u128) -> u64 {
        n as _
    }
}
impl DoubleInt for u32 {
    type Double = u64;

    fn as_double(self) -> u64 {
        self as _
    }
    fn from_double(n: u64) -> u32 {
        n as _
    }
}

// extended Euclid algorithm
// precondition: a is odd
pub(crate) fn modular_inverse<T: Int>(a: T) -> T {
    let zero = T::zero();
    let one = T::one();
    debug_assert!(a % (one + one) == one, "{:?} is not odd", a);

    let mut t = zero;
    let mut newt = one;
    let mut r = zero;
    let mut newr = a;

    while newr != zero {
        let quot = if r == zero {
            // special case when we're just starting out
            // This works because we know that
            // a does not divide 2^64, so floor(2^64 / a) == floor((2^64-1) / a);
            T::max_value()
        } else {
            r
        } / newr;

        let newtp = t.wrapping_sub(&quot.wrapping_mul(&newt));
        t = newt;
        newt = newtp;

        let newrp = r.wrapping_sub(&quot.wrapping_mul(&newr));
        r = newr;
        newr = newrp;
    }

    debug_assert_eq!(r, one);
    t
}

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

    fn test_inverter<T: Int>() {
        // All odd integers from 1 to 20 000
        let one = T::from(1).unwrap();
        let two = T::from(2).unwrap();
        let mut test_values = (0..10_000)
            .map(|i| T::from(i).unwrap())
            .map(|i| two * i + one);

        assert!(test_values.all(|x| x.wrapping_mul(&modular_inverse(x)) == one));
    }

    #[test]
    fn test_inverter_u32() {
        test_inverter::<u32>()
    }

    #[test]
    fn test_inverter_u64() {
        test_inverter::<u64>()
    }

    fn test_add<A: Arithmetic>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = A::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    println!("{n:?}, {x:?}, {y:?}", n = n, x = x, y = y);
                    assert_eq!((x + y) % n, m.to_u64(m.add(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_add_m32() {
        test_add::<Montgomery<u32>>()
    }

    #[test]
    fn test_add_m64() {
        test_add::<Montgomery<u64>>()
    }

    fn test_mult<A: Arithmetic>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = A::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    assert_eq!((x * y) % n, m.to_u64(m.mul(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_mult_m32() {
        test_mult::<Montgomery<u32>>()
    }

    #[test]
    fn test_mult_m64() {
        test_mult::<Montgomery<u64>>()
    }

    fn test_roundtrip<A: Arithmetic>() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = A::new(n);
            for x in 0..n {
                let x_ = m.from_u64(x);
                assert_eq!(x, m.to_u64(x_));
            }
        }
    }

    #[test]
    fn test_roundtrip_m32() {
        test_roundtrip::<Montgomery<u32>>()
    }

    #[test]
    fn test_roundtrip_m64() {
        test_roundtrip::<Montgomery<u64>>()
    }
}
<|endoftext|>


<commit_before>use crate::shapes::{FnInterp, Scratchpad};
use russell_lab::{mat_vec_mul, Matrix, Vector};

/// Calculates the real coordinates x from reference coordinates ξ
///
/// # Input
///
/// * `xxt` -- is the transposed coordinates matrix `Xᵀ` shown below (space_ndim,nnode)
///
/// The transposed coordinates matrix (real space) is such that:
///
/// ```text
///      ┌                              ┐  superscript = node
///      | x⁰₀  x¹₀  x²₀  x³₀       xᴹ₀ |  subscript = dimension
/// Xᵀ = | x⁰₁  x¹₁  x²₁  x³₁  ...  xᴹ₁ |
///      | x⁰₂  x¹₂  x²₂  x³₂       xᴹ₂ |
///      └                              ┘_(space_ndim,nnode)
/// ```
///
/// where `M = nnode - 1`
///
/// # Panics
///
/// This function does not check for the vector/matrix dimensions. Thus a panic may occur
/// if they are incompatible, including if they are not consistent with `fn_interp`.
#[inline]
pub fn calc_coords(x: &mut Vector, pad: &mut Scratchpad, ksi: &[f64], xxt: &Matrix, fn_interp: FnInterp) {
    fn_interp(&mut pad.interp, ksi);
    mat_vec_mul(x, 1.0, &xxt, &pad.interp).unwrap();
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#[cfg(test)]
mod tests {
    use super::calc_coords;
    use crate::shapes::op::testing::aux;
    use crate::shapes::{GeoKind, Scratchpad};
    use crate::util::ONE_BY_3;
    use crate::StrError;
    use russell_chk::assert_vec_approx_eq;
    use russell_lab::Vector;

    #[test]
    fn calc_coords_works() -> Result<(), StrError> {
        // kind, tol, tol_in for the case inside the shape
        let problem = vec![
            (GeoKind::Tri3, 1e-15, 0.45), // linear maps are inaccurate for the circular wedge
            (GeoKind::Tri6, 1e-15, 0.02), // << quadratic mapping is inaccurate as well
            (GeoKind::Tri10, 1e-14, 1e-14),
            (GeoKind::Tri15, 1e-14, 1e-4), // << this triangle is inaccurate as well here
            (GeoKind::Qua4, 1e-15, 0.14),  // linear maps are inaccurate for the circular wedge
            (GeoKind::Qua8, 1e-15, 1e-14),
            (GeoKind::Qua17, 1e-15, 1e-14),
            (GeoKind::Tet4, 1e-15, 0.45),   // linear tetrahedron is also inaccurate here
            (GeoKind::Tet10, 1e-15, 0.02),  // quadratic tetrahedron is also inaccurate here
            (GeoKind::Tet20, 1e-14, 1e-14), // cubic tetrahedron
            (GeoKind::Hex8, 1e-14, 0.14),   // bi-linear maps are inaccurate for the circular wedge
            (GeoKind::Hex20, 1e-15, 1e-14),
            (GeoKind::Hex32, 1e-15, 1e-4), // TODO: check why this tolerance is high
        ];

        // loop over shapes
        for (kind, tol, tol_in) in problem {
            // generate coordinates matrix
            let geo_ndim = kind.ndim();
            let space_ndim = usize::max(2, geo_ndim);
            let xxt = aux::gen_coords_transp(space_ndim, kind);

            // scratchpad and interpolation function
            let mut pad = Scratchpad::new(space_ndim, kind)?;
            let fn_interp = kind.functions().0;

            // loop over nodes of shape
            let nnode = kind.nnode();
            let mut x = Vector::new(space_ndim);
            let mut x_correct = Vector::new(space_ndim);
            for m in 0..nnode {
                // get ξᵐ corresponding to node m
                let ksi = kind.reference_coords(m);

                // calculate xᵐ(ξᵐ) using the isoparametric formula
                calc_coords(&mut x, &mut pad, ksi, &xxt, fn_interp);

                // compare xᵐ with generated coordinates
                aux::map_point_coords(&mut x_correct, ksi, kind);
                assert_vec_approx_eq!(x.as_data(), x_correct.as_data(), tol);
            }

            // test again inside the reference domain
            let ksi_in = if kind.is_tri_or_tet() {
                vec![ONE_BY_3; geo_ndim]
            } else {
                vec![0.0; geo_ndim]
            };
            calc_coords(&mut x, &mut pad, &ksi_in, &xxt, fn_interp);
            aux::map_point_coords(&mut x_correct, &ksi_in, kind);
            assert_vec_approx_eq!(x.as_data(), x_correct.as_data(), tol_in);
        }
        Ok(())
    }
}
<commit_msg>Improve doc comments<commit_after>use crate::shapes::{FnInterp, Scratchpad};
use russell_lab::{mat_vec_mul, Matrix, Vector};

/// Calculates the real coordinates x from reference coordinates ξ
///
/// This function uses the isoparametric formula to calculate x given ξ:
///
/// ```text
/// → →         →  →
/// x(ξ) = Σ Nᵐ(ξ) xᵐ
///        m
///
/// x := Xᵀ ⋅ N
/// ```
///
/// # Output
///
/// * `x` -- real coordinates (space_ndim)
/// * `pad.interp` -- (nnode) interpolation functions @ ξ
///
/// # Input
///
/// * `ksi` -- reference coordinates ξ with len ≥ geo_ndim
/// * `xxt` -- is the transposed coordinates matrix `Xᵀ` shown below (space_ndim,nnode)
///
/// The transposed coordinates matrix (real space) is such that:
///
/// ```text
///      ┌                              ┐  superscript = node
///      | x⁰₀  x¹₀  x²₀  x³₀       xᴹ₀ |  subscript = dimension
/// Xᵀ = | x⁰₁  x¹₁  x²₁  x³₁  ...  xᴹ₁ |
///      | x⁰₂  x¹₂  x²₂  x³₂       xᴹ₂ |
///      └                              ┘_(space_ndim,nnode)
/// ```
///
/// where `M = nnode - 1`
///
/// # Panics
///
/// This function does not check for the vector/matrix dimensions. Thus a panic may occur
/// if they are incompatible, including if they are not consistent with `fn_interp`.
#[inline]
pub fn calc_coords(x: &mut Vector, pad: &mut Scratchpad, ksi: &[f64], xxt: &Matrix, fn_interp: FnInterp) {
    fn_interp(&mut pad.interp, ksi);
    mat_vec_mul(x, 1.0, &xxt, &pad.interp).unwrap();
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#[cfg(test)]
mod tests {
    use super::calc_coords;
    use crate::shapes::op::testing::aux;
    use crate::shapes::{GeoKind, Scratchpad};
    use crate::util::ONE_BY_3;
    use crate::StrError;
    use russell_chk::assert_vec_approx_eq;
    use russell_lab::Vector;

    #[test]
    fn calc_coords_works() -> Result<(), StrError> {
        // kind, tol, tol_in for the case inside the shape
        let problem = vec![
            (GeoKind::Tri3, 1e-15, 0.45), // linear maps are inaccurate for the circular wedge
            (GeoKind::Tri6, 1e-15, 0.02), // << quadratic mapping is inaccurate as well
            (GeoKind::Tri10, 1e-14, 1e-14),
            (GeoKind::Tri15, 1e-14, 1e-4), // << this triangle is inaccurate as well here
            (GeoKind::Qua4, 1e-15, 0.14),  // linear maps are inaccurate for the circular wedge
            (GeoKind::Qua8, 1e-15, 1e-14),
            (GeoKind::Qua17, 1e-15, 1e-14),
            (GeoKind::Tet4, 1e-15, 0.45),   // linear tetrahedron is also inaccurate here
            (GeoKind::Tet10, 1e-15, 0.02),  // quadratic tetrahedron is also inaccurate here
            (GeoKind::Tet20, 1e-14, 1e-14), // cubic tetrahedron
            (GeoKind::Hex8, 1e-14, 0.14),   // bi-linear maps are inaccurate for the circular wedge
            (GeoKind::Hex20, 1e-15, 1e-14),
            (GeoKind::Hex32, 1e-15, 1e-4), // TODO: check why this tolerance is high
        ];

        // loop over shapes
        for (kind, tol, tol_in) in problem {
            // generate coordinates matrix
            let geo_ndim = kind.ndim();
            let space_ndim = usize::max(2, geo_ndim);
            let xxt = aux::gen_coords_transp(space_ndim, kind);

            // scratchpad and interpolation function
            let mut pad = Scratchpad::new(space_ndim, kind)?;
            let fn_interp = kind.functions().0;

            // loop over nodes of shape
            let nnode = kind.nnode();
            let mut x = Vector::new(space_ndim);
            let mut x_correct = Vector::new(space_ndim);
            for m in 0..nnode {
                // get ξᵐ corresponding to node m
                let ksi = kind.reference_coords(m);

                // calculate xᵐ(ξᵐ) using the isoparametric formula
                calc_coords(&mut x, &mut pad, ksi, &xxt, fn_interp);

                // compare xᵐ with generated coordinates
                aux::map_point_coords(&mut x_correct, ksi, kind);
                assert_vec_approx_eq!(x.as_data(), x_correct.as_data(), tol);
            }

            // test again inside the reference domain
            let ksi_in = if kind.is_tri_or_tet() {
                vec![ONE_BY_3; geo_ndim]
            } else {
                vec![0.0; geo_ndim]
            };
            calc_coords(&mut x, &mut pad, &ksi_in, &xxt, fn_interp);
            aux::map_point_coords(&mut x_correct, &ksi_in, kind);
            assert_vec_approx_eq!(x.as_data(), x_correct.as_data(), tol_in);
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>use core::mem;
use core::ptr::Unique;

use crate::memory::{allocate_frames, deallocate_frames, Frame};

use super::{ActivePageTable, Page, PAGE_SIZE, PhysicalAddress, VirtualAddress};
use super::entry::EntryFlags;
use super::table::{self, Table, Level4};

/// In order to enforce correct paging operations in the kernel, these types
/// are returned on any mapping operation to get the code involved to specify
/// how it intends to flush changes to a page table
#[must_use = "The page table must be flushed, or the changes unsafely ignored"]
pub struct MapperFlush(Page);

impl MapperFlush {
    /// Create a new page flush promise
    pub fn new(page: Page) -> MapperFlush {
        MapperFlush(page)
    }

    /// Flush this page in the active table
    pub fn flush(self, table: &mut ActivePageTable) {
        table.flush(self.0);
        mem::forget(self);
    }

    /// Ignore the flush. This is unsafe, and a reason should be provided for use
    pub unsafe fn ignore(self) {
        mem::forget(self);
    }
}

/// A flush cannot be dropped, it must be consumed
impl Drop for MapperFlush {
    fn drop(&mut self) {
        panic!("Mapper flush was not utilized");
    }
}

/// To allow for combining multiple flushes into one, we have a way of flushing
/// the active table, which can consume `MapperFlush` structs
#[must_use = "The page table must be flushed, or the changes unsafely ignored"]
pub struct MapperFlushAll(bool);

impl MapperFlushAll {
    /// Create a new promise to flush all mappings
    pub fn new() -> MapperFlushAll {
        MapperFlushAll(false)
    }

    /// Consume a single page flush
    pub fn consume(&mut self, flush: MapperFlush) {
        self.0 = true;
        mem::forget(flush);
    }

    /// Flush the active page table
    pub fn flush(self, table: &mut ActivePageTable) {
        if self.0 {
            table.flush_all();
        }
        mem::forget(self);
    }

    /// Ignore the flush. This is unsafe, and a reason should be provided for use
    pub unsafe fn ignore(self) {
        mem::forget(self);
    }
}

/// A flush cannot be dropped, it must be consumed
impl Drop for MapperFlushAll {
    fn drop(&mut self) {
        panic!("Mapper flush all was not utilized");
    }
}

pub struct Mapper {
    p4: Unique<Table<Level4>>,
}

impl Mapper {
    /// Create a new page table
    pub unsafe fn new() -> Mapper {
        Mapper {
            p4: Unique::new_unchecked(table::P4),
        }
    }

    pub fn p4(&self) -> &Table<Level4> {
        unsafe { self.p4.as_ref() }
    }

    pub fn p4_mut(&mut self) -> &mut Table<Level4> {
        unsafe { self.p4.as_mut() }
    }

    /// Map a page to a frame
    pub fn map_to(&mut self, page: Page, frame: Frame, flags: EntryFlags) -> MapperFlush {
        let p3 = self.p4_mut().next_table_create(page.p4_index());
        let p2 = p3.next_table_create(page.p3_index());
        let p1 = p2.next_table_create(page.p2_index());

        assert!(p1[page.p1_index()].is_unused(),
            "{:X}: Set to {:X}: {:?}, requesting {:X}: {:?}",
            page.start_address().get(),
            p1[page.p1_index()].address().get(), p1[page.p1_index()].flags(),
            frame.start_address().get(), flags);
        p1.increment_entry_count();
        p1[page.p1_index()].set(frame, flags | EntryFlags::PRESENT);
        MapperFlush::new(page)
    }

    /// Map a page to the next free frame
    pub fn map(&mut self, page: Page, flags: EntryFlags) -> MapperFlush {
        let frame = allocate_frames(1).expect("out of frames");
        self.map_to(page, frame, flags)
    }

    /// Update flags for a page
    pub fn remap(&mut self, page: Page, flags: EntryFlags) -> MapperFlush {
        let p3 = self.p4_mut().next_table_mut(page.p4_index()).expect("failed to remap: no p3");
        let p2 = p3.next_table_mut(page.p3_index()).expect("failed to remap: no p2");
        let p1 = p2.next_table_mut(page.p2_index()).expect("failed to remap: no p1");
        let frame = p1[page.p1_index()].pointed_frame().expect("failed to remap: not mapped");
        p1[page.p1_index()].set(frame, flags | EntryFlags::PRESENT);
        MapperFlush::new(page)
    }

    /// Identity map a frame
    pub fn identity_map(&mut self, frame: Frame, flags: EntryFlags) -> MapperFlush {
        let page = Page::containing_address(VirtualAddress::new(frame.start_address().get()));
        self.map_to(page, frame, flags)
    }

    fn unmap_inner(&mut self, page: Page, keep_parents: bool) -> Frame {
        let frame;

        let p4 = self.p4_mut();
        if let Some(p3) = p4.next_table_mut(page.p4_index()) {
            if let Some(p2) = p3.next_table_mut(page.p3_index()) {
                if let Some(p1) = p2.next_table_mut(page.p2_index()) {
                    frame = if let Some(frame) = p1[page.p1_index()].pointed_frame() {
                        frame
                    } else {
                        panic!("unmap_inner({:X}): frame not found", page.start_address().get())
                    };

                    p1.decrement_entry_count();
                    p1[page.p1_index()].set_unused();

                    if keep_parents || ! p1.is_unused() {
                        return frame;
                    }
                } else {
                    panic!("unmap_inner({:X}): p1 not found", page.start_address().get());
                }

                if let Some(p1_frame) = p2[page.p2_index()].pointed_frame() {
                    //println!("Free p1 {:?}", p1_frame);
                    p2.decrement_entry_count();
                    p2[page.p2_index()].set_unused();
                    deallocate_frames(p1_frame, 1);
                } else {
                    panic!("unmap_inner({:X}): p1_frame not found", page.start_address().get());
                }

                if ! p2.is_unused() {
                    return frame;
                }
            } else {
                panic!("unmap_inner({:X}): p2 not found", page.start_address().get());
            }

            if let Some(p2_frame) = p3[page.p3_index()].pointed_frame() {
                //println!("Free p2 {:?}", p2_frame);
                p3.decrement_entry_count();
                p3[page.p3_index()].set_unused();
                deallocate_frames(p2_frame, 1);
            } else {
                panic!("unmap_inner({:X}): p2_frame not found", page.start_address().get());
            }

            if ! p3.is_unused() {
                return frame;
            }
        } else {
            panic!("unmap_inner({:X}): p3 not found", page.start_address().get());
        }

        if let Some(p3_frame) = p4[page.p4_index()].pointed_frame() {
            //println!("Free p3 {:?}", p3_frame);
            p4.decrement_entry_count();
            p4[page.p4_index()].set_unused();
            deallocate_frames(p3_frame, 1);
        } else {
            panic!("unmap_inner({:X}): p3_frame not found", page.start_address().get());
        }

        frame
    }

    /// Unmap a page
    pub fn unmap(&mut self, page: Page) -> MapperFlush {
        let frame = self.unmap_inner(page, false);
        deallocate_frames(frame, 1);
        MapperFlush::new(page)
    }

    /// Unmap a page, return frame without free
    pub fn unmap_return(&mut self, page: Page, keep_parents: bool) -> (MapperFlush, Frame) {
        let frame = self.unmap_inner(page, keep_parents);
        (MapperFlush::new(page), frame)
    }

    pub fn translate_page(&self, page: Page) -> Option<Frame> {
        self.p4().next_table(page.p4_index())
            .and_then(|p3| p3.next_table(page.p3_index()))
            .and_then(|p2| p2.next_table(page.p2_index()))
            .and_then(|p1| p1[page.p1_index()].pointed_frame())
    }

    pub fn translate_page_flags(&self, page: Page) -> Option<EntryFlags> {
        self.p4().next_table(page.p4_index())
            .and_then(|p3| p3.next_table(page.p3_index()))
            .and_then(|p2| p2.next_table(page.p2_index()))
            .and_then(|p1| Some(p1[page.p1_index()].flags()))
    }

    /// Translate a virtual address to a physical one
    pub fn translate(&self, virtual_address: VirtualAddress) -> Option<PhysicalAddress> {
        let offset = virtual_address.get() % PAGE_SIZE;
        self.translate_page(Page::containing_address(virtual_address))
            .map(|frame| PhysicalAddress::new(frame.start_address().get() + offset))
    }
}
<commit_msg>Work around unmapping user heap<commit_after>use core::mem;
use core::ptr::Unique;

use crate::memory::{allocate_frames, deallocate_frames, Frame};

use super::{ActivePageTable, Page, PAGE_SIZE, PhysicalAddress, VirtualAddress};
use super::entry::EntryFlags;
use super::table::{self, Table, Level4};

/// In order to enforce correct paging operations in the kernel, these types
/// are returned on any mapping operation to get the code involved to specify
/// how it intends to flush changes to a page table
#[must_use = "The page table must be flushed, or the changes unsafely ignored"]
pub struct MapperFlush(Page);

impl MapperFlush {
    /// Create a new page flush promise
    pub fn new(page: Page) -> MapperFlush {
        MapperFlush(page)
    }

    /// Flush this page in the active table
    pub fn flush(self, table: &mut ActivePageTable) {
        table.flush(self.0);
        mem::forget(self);
    }

    /// Ignore the flush. This is unsafe, and a reason should be provided for use
    pub unsafe fn ignore(self) {
        mem::forget(self);
    }
}

/// A flush cannot be dropped, it must be consumed
impl Drop for MapperFlush {
    fn drop(&mut self) {
        panic!("Mapper flush was not utilized");
    }
}

/// To allow for combining multiple flushes into one, we have a way of flushing
/// the active table, which can consume `MapperFlush` structs
#[must_use = "The page table must be flushed, or the changes unsafely ignored"]
pub struct MapperFlushAll(bool);

impl MapperFlushAll {
    /// Create a new promise to flush all mappings
    pub fn new() -> MapperFlushAll {
        MapperFlushAll(false)
    }

    /// Consume a single page flush
    pub fn consume(&mut self, flush: MapperFlush) {
        self.0 = true;
        mem::forget(flush);
    }

    /// Flush the active page table
    pub fn flush(self, table: &mut ActivePageTable) {
        if self.0 {
            table.flush_all();
        }
        mem::forget(self);
    }

    /// Ignore the flush. This is unsafe, and a reason should be provided for use
    pub unsafe fn ignore(self) {
        mem::forget(self);
    }
}

/// A flush cannot be dropped, it must be consumed
impl Drop for MapperFlushAll {
    fn drop(&mut self) {
        panic!("Mapper flush all was not utilized");
    }
}

pub struct Mapper {
    p4: Unique<Table<Level4>>,
}

impl Mapper {
    /// Create a new page table
    pub unsafe fn new() -> Mapper {
        Mapper {
            p4: Unique::new_unchecked(table::P4),
        }
    }

    pub fn p4(&self) -> &Table<Level4> {
        unsafe { self.p4.as_ref() }
    }

    pub fn p4_mut(&mut self) -> &mut Table<Level4> {
        unsafe { self.p4.as_mut() }
    }

    /// Map a page to a frame
    pub fn map_to(&mut self, page: Page, frame: Frame, flags: EntryFlags) -> MapperFlush {
        let p3 = self.p4_mut().next_table_create(page.p4_index());
        let p2 = p3.next_table_create(page.p3_index());
        let p1 = p2.next_table_create(page.p2_index());

        assert!(p1[page.p1_index()].is_unused(),
            "{:X}: Set to {:X}: {:?}, requesting {:X}: {:?}",
            page.start_address().get(),
            p1[page.p1_index()].address().get(), p1[page.p1_index()].flags(),
            frame.start_address().get(), flags);
        p1.increment_entry_count();
        p1[page.p1_index()].set(frame, flags | EntryFlags::PRESENT);
        MapperFlush::new(page)
    }

    /// Map a page to the next free frame
    pub fn map(&mut self, page: Page, flags: EntryFlags) -> MapperFlush {
        let frame = allocate_frames(1).expect("out of frames");
        self.map_to(page, frame, flags)
    }

    /// Update flags for a page
    pub fn remap(&mut self, page: Page, flags: EntryFlags) -> MapperFlush {
        let p3 = self.p4_mut().next_table_mut(page.p4_index()).expect("failed to remap: no p3");
        let p2 = p3.next_table_mut(page.p3_index()).expect("failed to remap: no p2");
        let p1 = p2.next_table_mut(page.p2_index()).expect("failed to remap: no p1");
        let frame = p1[page.p1_index()].pointed_frame().expect("failed to remap: not mapped");
        p1[page.p1_index()].set(frame, flags | EntryFlags::PRESENT);
        MapperFlush::new(page)
    }

    /// Identity map a frame
    pub fn identity_map(&mut self, frame: Frame, flags: EntryFlags) -> MapperFlush {
        let page = Page::containing_address(VirtualAddress::new(frame.start_address().get()));
        self.map_to(page, frame, flags)
    }

    fn unmap_inner(&mut self, page: Page, keep_parents: bool) -> Frame {
        let frame;

        let p4 = self.p4_mut();
        if let Some(p3) = p4.next_table_mut(page.p4_index()) {
            if let Some(p2) = p3.next_table_mut(page.p3_index()) {
                if let Some(p1) = p2.next_table_mut(page.p2_index()) {
                    frame = if let Some(frame) = p1[page.p1_index()].pointed_frame() {
                        frame
                    } else {
                        panic!("unmap_inner({:X}): frame not found", page.start_address().get())
                    };

                    p1.decrement_entry_count();
                    p1[page.p1_index()].set_unused();

                    if keep_parents || ! p1.is_unused() {
                        return frame;
                    }
                } else {
                    panic!("unmap_inner({:X}): p1 not found", page.start_address().get());
                }

                if let Some(p1_frame) = p2[page.p2_index()].pointed_frame() {
                    //println!("Free p1 {:?}", p1_frame);
                    p2.decrement_entry_count();
                    p2[page.p2_index()].set_unused();
                    deallocate_frames(p1_frame, 1);
                } else {
                    panic!("unmap_inner({:X}): p1_frame not found", page.start_address().get());
                }

                if ! p2.is_unused() {
                    return frame;
                }
            } else {
                panic!("unmap_inner({:X}): p2 not found", page.start_address().get());
            }

            if let Some(p2_frame) = p3[page.p3_index()].pointed_frame() {
                //println!("Free p2 {:?}", p2_frame);
                p3.decrement_entry_count();
                p3[page.p3_index()].set_unused();
                deallocate_frames(p2_frame, 1);
            } else {
                panic!("unmap_inner({:X}): p2_frame not found", page.start_address().get());
            }

            if ! p3.is_unused() {
                return frame;
            }
        } else {
            panic!("unmap_inner({:X}): p3 not found", page.start_address().get());
        }

        if let Some(p3_frame) = p4[page.p4_index()].pointed_frame() {
            //TODO: Find out why this breaks user heap
            //println!("Free p3 {:?}", p3_frame);
            //p4.decrement_entry_count();
            //p4[page.p4_index()].set_unused();
            //deallocate_frames(p3_frame, 1);
        } else {
            panic!("unmap_inner({:X}): p3_frame not found", page.start_address().get());
        }

        frame
    }

    /// Unmap a page
    pub fn unmap(&mut self, page: Page) -> MapperFlush {
        let frame = self.unmap_inner(page, false);
        deallocate_frames(frame, 1);
        MapperFlush::new(page)
    }

    /// Unmap a page, return frame without free
    pub fn unmap_return(&mut self, page: Page, keep_parents: bool) -> (MapperFlush, Frame) {
        let frame = self.unmap_inner(page, keep_parents);
        (MapperFlush::new(page), frame)
    }

    pub fn translate_page(&self, page: Page) -> Option<Frame> {
        self.p4().next_table(page.p4_index())
            .and_then(|p3| p3.next_table(page.p3_index()))
            .and_then(|p2| p2.next_table(page.p2_index()))
            .and_then(|p1| p1[page.p1_index()].pointed_frame())
    }

    pub fn translate_page_flags(&self, page: Page) -> Option<EntryFlags> {
        self.p4().next_table(page.p4_index())
            .and_then(|p3| p3.next_table(page.p3_index()))
            .and_then(|p2| p2.next_table(page.p2_index()))
            .and_then(|p1| Some(p1[page.p1_index()].flags()))
    }

    /// Translate a virtual address to a physical one
    pub fn translate(&self, virtual_address: VirtualAddress) -> Option<PhysicalAddress> {
        let offset = virtual_address.get() % PAGE_SIZE;
        self.translate_page(Page::containing_address(virtual_address))
            .map(|frame| PhysicalAddress::new(frame.start_address().get() + offset))
    }
}
<|endoftext|>


<commit_before>#![allow(clippy::nonstandard_macro_braces)] // clippy bug, see https://github.com/rust-lang/rust-clippy/issues/7434

use anyhow::Result;
use bitflags::bitflags;
use thiserror::Error;

use crate::{
    arch::Arch,
    aspace::{AbsoluteAddressSpace, AddressSpace},
    pagemap::PageMapError::NotMapped,
    RVA, VA,
};

#[derive(Error, Debug)]
pub enum ModuleError {
    #[error("invalid address: {0:#x}")]
    InvalidAddress(u64),
}

bitflags! {
    pub struct Permissions: u8 {
        const R = 0b0000_0001;
        const W = 0b0000_0010;
        const X = 0b0000_0100;
        const RW = Self::R.bits | Self::W.bits;
        const RX =  Self::R.bits | Self::X.bits;
        const WX =  Self::W.bits | Self::X.bits;
        const RWX =  Self::R.bits | Self::W.bits | Self::X.bits;
    }
}

#[derive(Debug, Clone)]
pub struct Section {
    // source data, from the PE file, relative to file start.
    pub physical_range: std::ops::Range<RVA>,
    // as mapped into memory with absolute addresses.
    pub virtual_range:  std::ops::Range<VA>,
    pub permissions:    Permissions,
    pub name:           String,
}

/// An address space, as a file would be loaded into memory.
/// This has an associated architecture (e.g. x32 or x64),
/// base address, and collection of sections.
/// This is the information that we'd expect to be common across formats.
#[derive(Clone)]
pub struct Module {
    pub arch:          Arch,
    pub sections:      Vec<Section>,
    pub address_space: AbsoluteAddressSpace,
}

impl Module {
    pub fn read_va_at_rva(&self, offset: RVA) -> Result<VA> {
        match self.arch {
            Arch::X32 => Ok(self.address_space.relative.read_u32(offset)? as VA),
            Arch::X64 => Ok(self.address_space.relative.read_u64(offset)? as VA),
        }
    }

    pub fn read_rva_at_rva(&self, offset: RVA) -> Result<RVA> {
        match self.arch {
            Arch::X32 => Ok(self.address_space.relative.read_u32(offset)? as RVA),
            Arch::X64 => Ok(self.address_space.relative.read_u64(offset)? as RVA),
        }
    }

    pub fn read_va_at_va(&self, offset: VA) -> Result<VA> {
        match self.arch {
            Arch::X32 => Ok(self.address_space.read_u32(offset)? as VA),
            Arch::X64 => Ok(self.address_space.read_u64(offset)? as VA),
        }
    }

    pub fn read_rva_at_va(&self, offset: VA) -> Result<RVA> {
        match self.arch {
            Arch::X32 => Ok(self.address_space.read_u32(offset)? as RVA),
            Arch::X64 => Ok(self.address_space.read_u64(offset)? as RVA),
        }
    }

    pub fn probe_va(&self, offset: VA, perm: Permissions) -> bool {
        self.sections
            .iter()
            .any(|section| section.virtual_range.contains(&offset) && section.permissions.intersects(perm))
    }

    pub fn probe_rva(&self, offset: RVA, perm: Permissions) -> bool {
        let va = self.address_space.base_address + offset;
        self.probe_va(va, perm)
    }

    pub fn file_offset(&self, va: VA) -> Result<usize> {
        if let Some(sec) = self.sections.iter().find(|&sec| sec.virtual_range.contains(&va)) {
            let offset = va - sec.virtual_range.start;
            Ok((offset + sec.physical_range.start) as usize)
        } else {
            Err(NotMapped.into())
        }
    }

    pub fn virtual_address(&self, file_offset: u64) -> Result<VA> {
        self.sections
            .iter()
            .find(|&sec| sec.physical_range.contains(&(file_offset as u64)))
            .map(|sec| {
                let section_offset = file_offset as u64 - sec.physical_range.start;
                sec.virtual_range.start + section_offset
            })
            .ok_or_else(|| NotMapped.into())
    }
}
<commit_msg>module: permissions: add Display<commit_after>#![allow(clippy::nonstandard_macro_braces)] // clippy bug, see https://github.com/rust-lang/rust-clippy/issues/7434

use anyhow::Result;
use bitflags::bitflags;
use thiserror::Error;

use crate::{
    arch::Arch,
    aspace::{AbsoluteAddressSpace, AddressSpace},
    pagemap::PageMapError::NotMapped,
    RVA, VA,
};

#[derive(Error, Debug)]
pub enum ModuleError {
    #[error("invalid address: {0:#x}")]
    InvalidAddress(u64),
}

bitflags! {
    pub struct Permissions: u8 {
        const R = 0b0000_0001;
        const W = 0b0000_0010;
        const X = 0b0000_0100;
        const RW = Self::R.bits | Self::W.bits;
        const RX =  Self::R.bits | Self::X.bits;
        const WX =  Self::W.bits | Self::X.bits;
        const RWX =  Self::R.bits | Self::W.bits | Self::X.bits;
    }
}

impl std::fmt::Display for Permissions {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        if self.intersects(Permissions::R) {
            write!(f, "R")?;
        }
        if self.intersects(Permissions::W) {
            write!(f, "W")?;
        }
        if self.intersects(Permissions::X) {
            write!(f, "X")?;
        }
        Ok(())
    }
}

#[derive(Debug, Clone)]
pub struct Section {
    // source data, from the PE file, relative to file start.
    pub physical_range: std::ops::Range<RVA>,
    // as mapped into memory with absolute addresses.
    pub virtual_range:  std::ops::Range<VA>,
    pub permissions:    Permissions,
    pub name:           String,
}

/// An address space, as a file would be loaded into memory.
/// This has an associated architecture (e.g. x32 or x64),
/// base address, and collection of sections.
/// This is the information that we'd expect to be common across formats.
#[derive(Clone)]
pub struct Module {
    pub arch:          Arch,
    pub sections:      Vec<Section>,
    pub address_space: AbsoluteAddressSpace,
}

impl Module {
    pub fn read_va_at_rva(&self, offset: RVA) -> Result<VA> {
        match self.arch {
            Arch::X32 => Ok(self.address_space.relative.read_u32(offset)? as VA),
            Arch::X64 => Ok(self.address_space.relative.read_u64(offset)? as VA),
        }
    }

    pub fn read_rva_at_rva(&self, offset: RVA) -> Result<RVA> {
        match self.arch {
            Arch::X32 => Ok(self.address_space.relative.read_u32(offset)? as RVA),
            Arch::X64 => Ok(self.address_space.relative.read_u64(offset)? as RVA),
        }
    }

    pub fn read_va_at_va(&self, offset: VA) -> Result<VA> {
        match self.arch {
            Arch::X32 => Ok(self.address_space.read_u32(offset)? as VA),
            Arch::X64 => Ok(self.address_space.read_u64(offset)? as VA),
        }
    }

    pub fn read_rva_at_va(&self, offset: VA) -> Result<RVA> {
        match self.arch {
            Arch::X32 => Ok(self.address_space.read_u32(offset)? as RVA),
            Arch::X64 => Ok(self.address_space.read_u64(offset)? as RVA),
        }
    }

    pub fn probe_va(&self, offset: VA, perm: Permissions) -> bool {
        self.sections
            .iter()
            .any(|section| section.virtual_range.contains(&offset) && section.permissions.intersects(perm))
    }

    pub fn probe_rva(&self, offset: RVA, perm: Permissions) -> bool {
        let va = self.address_space.base_address + offset;
        self.probe_va(va, perm)
    }

    pub fn file_offset(&self, va: VA) -> Result<usize> {
        if let Some(sec) = self.sections.iter().find(|&sec| sec.virtual_range.contains(&va)) {
            let offset = va - sec.virtual_range.start;
            Ok((offset + sec.physical_range.start) as usize)
        } else {
            Err(NotMapped.into())
        }
    }

    pub fn virtual_address(&self, file_offset: u64) -> Result<VA> {
        self.sections
            .iter()
            .find(|&sec| sec.physical_range.contains(&(file_offset as u64)))
            .map(|sec| {
                let section_offset = file_offset as u64 - sec.physical_range.start;
                sec.virtual_range.start + section_offset
            })
            .ok_or_else(|| NotMapped.into())
    }
}
<|endoftext|>


<commit_before>//
// gash.rs
//
// Starting code for PS2
// Running on Rust 0.9
//
// University of Virginia - cs4414 Spring 2014
// Weilin Xu, David Evans
// Version 0.4
//

extern mod extra;

use std::{io, run, os};
use std::io::buffered::BufferedReader;
// use std::io::stdio::StdReader;
// use std::io::Reader;
// use std::io::Writer;
use std::io::stdin;
use extra::getopts;

struct Shell {
    cmd_prompt: ~str,
    cmd_history: ~[~str],
}

impl Shell {
    fn new(prompt_str: &str) -> Shell {
        Shell {
            cmd_prompt: prompt_str.to_owned(),
            cmd_history: ~[],
        }
    }

    fn run(&mut self) {
        let mut stdin = BufferedReader::new(stdin());

        loop {
            print(self.cmd_prompt);
            io::stdio::flush();

            let mut line : ~str = stdin.read_line().unwrap().to_owned();
            let mut cmd_line: ~str = line.trim().to_owned();
            let mut background: bool = false;
            let mut i : uint = 0;

            // This block handles if there are no spaces around pipe.
            let lineClone = line.clone();
            for character in lineClone.chars() {
                if character == '<' {
                    line = line.slice(0, i).to_owned() + " < " + line.slice(i+1, line.char_len()).to_owned();
                    i = i + 2;
                } else if character == '>' {
                    line = line.slice(0, i).to_owned() + " > " + line.slice(i+1, line.char_len()).to_owned();
                    i = i + 2;
                } else if character == '|' {
                    line = line.slice(0, i).to_owned() + " | " + line.slice(i+1, line.char_len()).to_owned();
                    i = i + 2;
                }
                i = i + 1;
            }

            if cmd_line.char_len() > 0 {
                background = cmd_line.char_at(cmd_line.char_len() - 1) == '&';
                if background {
                    cmd_line = cmd_line.slice(0, cmd_line.char_len() - 1).trim().to_owned();
                }
            }

            let params = cmd_line.clone().to_owned();
            let program = cmd_line.splitn(' ', 1).nth(0).expect("no program");

            match program {
                ""          =>  { continue; }
                "exit"      =>  { return; }
                "cd"        =>  { Shell::run_check_mode(background, Shell::run_cd(params)); }
                "history"   =>  { Shell::run_check_mode(background, Shell::run_history(params, self.cmd_history)); }
                _           =>  { Shell::run_check_mode(background, Shell::run_cmdline(params)); }
            }

            self.cmd_history.push(program.to_owned());
        }
    }

    fn run_check_mode(background: bool, f: proc()) {
        if background {
            spawn(proc() { f(); });
        } else {
            f();
        }
    }

    fn run_cmdline_single(cmd_line: &str) {
        (Shell::run_cmdline(cmd_line))();
    }

    fn run_cmdline(cmd_line: &str) -> proc() {
        let params = cmd_line.to_owned();
        return proc() {
            let mut argv: ~[~str] = Shell::get_args(params);

            if argv.len() > 0 {
                let program: ~str = argv.remove(0);
                Shell::run_cmd(program, argv/*, None, false*/);
            }
        };
    }

    fn run_cmd(program: &str, argv: &[~str]/*, inputStream: Option<~Reader>, hasOutRedirect: bool*/) /*-> Option<~[u8]>*/ {
        if Shell::cmd_exists(program) {
            run::process_status(program, argv);

            // let mut options = run::ProcessOptions::new();
            // options.in_fd = match inputStream {
            //     Some(reader) => { None }
            //     None => { Some(0) }
            // };

            // if hasOutRedirect {
            //     options.out_fd = None;
            // } else {
            //     options.out_fd = Some(1);
            // }

            // let mut process = run::Process::new(program, argv, options).unwrap();
            // match options.in_fd {
            //     Some(_) => {  }
            //     None => {
            //         let buf = inputStream.read_to_str().into_bytes();
            //         process.input().write(buf);
            //     }
            // }

            // if hasOutRedirect {
            //     let mut processOutput = process.finish_with_output();
            //     return Some(processOutput.output);

            // } else {
            //     process.finish();
            //     return None;
            // }
        } else {
            println!("{:s}: command not found", program);
            // return None;
        }
    }

    fn cmd_exists(cmd_path: &str) -> bool {
        let ret = run::process_output("which", [cmd_path.to_owned()]);
        return ret.expect("exit code error.").status.success();
    }

    fn run_cd(cmd_line: &str) -> proc() {
        let params = cmd_line.to_owned();
        return proc() {
            let argv: ~[~str] = Shell::get_args(params);
            let pathOpt: Option<Path> = match argv.len() {
                1   =>  { os::homedir() }
                0   =>  { os::homedir() }
                _   =>  { Some(Path::new(argv[1])) }
            };

            match pathOpt {
                Some(path)   =>  {
                    if path.is_dir() {
                        os::change_dir(&path);
                    } else {
                        println!("Error: {:s} is not a directory", path.as_str().unwrap());
                    }
                }
                None        =>  {
                    println!("Error: Invalid path");
                }
            };
        };
    }

    fn run_history(cmd_line: &str, cmd_history: &[~str]) -> proc() {
        let history = cmd_history.to_owned();
        let params = cmd_line.to_owned();
        return proc() {
            let argv: ~[~str] = Shell::get_args(params);
            if (argv.len() > 1) {
                println!("Error: history does not take options (sadly).");
            }
            else {
                let mut i = 1;
                for entry in history.iter() {
                    println!("{:d} \t{:s}", i, entry.to_owned());
                    i = i + 1;
                }
            }
        };
    }

    fn get_args(cmd_line: &str) -> ~[~str] {
        return cmd_line.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();
    }
}

fn get_cmdline_from_args() -> Option<~str> {
    /* Begin processing program arguments and initiate the parameters. */
    let args = os::args();

    let opts = ~[
        getopts::optopt("c")
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m }
        Err(f) => { fail!(f.to_err_msg()) }
    };

    if matches.opt_present("c") {
        let cmd_str = match matches.opt_str("c") {
                                                Some(cmd_str) => {cmd_str.to_owned()},
                                                None => {~""}
                                              };
        return Some(cmd_str);
    } else {
        return None;
    }
}

fn main() {
    let opt_cmd_line = get_cmdline_from_args();

    match opt_cmd_line {
        Some(cmd_line) => Shell::run_cmdline_single(cmd_line),
        None           => Shell::new("gash > ").run()
    }
}
<commit_msg>added support for redirects and pipes (HALLELUJAH)<commit_after>//
// gash.rs
//
// Starting code for PS2
// Running on Rust 0.9
//
// University of Virginia - cs4414 Spring 2014
// Weilin Xu, David Evans
// Version 0.4
//

extern mod extra;

use std::{io, run, os, str};
use std::io::buffered::BufferedReader;
use std::path::posix::Path;
use std::io::fs::File;
use std::io::stdin;
use extra::getopts;

struct Shell {
    cmd_prompt: ~str,
    cmd_history: ~[~str],
}

impl Shell {
    fn new(prompt_str: &str) -> Shell {
        Shell {
            cmd_prompt: prompt_str.to_owned(),
            cmd_history: ~[],
        }
    }

    fn run(&mut self) {
        let mut stdin = BufferedReader::new(stdin());

        loop {
            print(self.cmd_prompt);
            io::stdio::flush();

            let mut line : ~str = stdin.read_line().unwrap().to_owned();
            let mut cmd_line: ~str = line.trim().to_owned();
            let mut background: bool = false;
            let mut i : uint = 0;

            // This block handles if there are no spaces around pipe.
            let lineClone = line.clone();
            for character in lineClone.chars() {
                if character == '<' {
                    line = line.slice(0, i).to_owned() + " < " + line.slice(i+1, line.char_len()).to_owned();
                    i = i + 2;
                } else if character == '>' {
                    line = line.slice(0, i).to_owned() + " > " + line.slice(i+1, line.char_len()).to_owned();
                    i = i + 2;
                } else if character == '|' {
                    line = line.slice(0, i).to_owned() + " | " + line.slice(i+1, line.char_len()).to_owned();
                    i = i + 2;
                }
                i = i + 1;
            }

            if cmd_line.char_len() > 0 {
                background = cmd_line.char_at(cmd_line.char_len() - 1) == '&';
                if background {
                    cmd_line = cmd_line.slice(0, cmd_line.char_len() - 1).trim().to_owned();
                }
            }

            let params = cmd_line.clone().to_owned();
            let program = cmd_line.splitn(' ', 1).nth(0).expect("no program");

            match program {
                ""          =>  { continue; }
                "exit"      =>  { return; }
                "cd"        =>  { Shell::run_check_mode(background, Shell::run_cd(params)); }
                "history"   =>  { Shell::run_check_mode(background, Shell::run_history(params, self.cmd_history)); }
                _           =>  { Shell::run_check_mode(background, Shell::run_cmdline(params)); }
            }

            self.cmd_history.push(program.to_owned());
        }
    }

    fn run_check_mode(background: bool, f: proc()) {
        if background {
            spawn(proc() { f(); });
        } else {
            f();
        }
    }

    fn run_cmdline_single(cmd_line: &str) {
        (Shell::run_cmdline(cmd_line))();
    }

    fn run_cmdline(cmd_line: &str) -> proc() {
        let params = cmd_line.to_owned();
        return proc() {
            if params.trim() != "" {
                let output = Shell::handle_pipes(params, false);
                // write output to file (if necessary and exists)
            }
        };
    }

    fn handle_pipes(cmd_line: &str, out_pipe: bool) -> Option<~str> {
        let pipeSplit = Shell::split_on_last_pipe(cmd_line);
        let mut argv = Shell::get_args_no_redirects(pipeSplit[0]);
        let program = argv.remove(0);
        let hasOutRedirect = match Shell::get_output_file(cmd_line) {
            Some(_) => { true }
            None => {
                if out_pipe {
                    true
                } else {
                    false
                }
            }
        };

        if pipeSplit.len() == 2 {
            return Shell::run_cmd(program, argv, Shell::handle_pipes(pipeSplit[1], true), hasOutRedirect);
        } else {
            return Shell::run_cmd(program, argv, Shell::get_input_file_contents(cmd_line), hasOutRedirect);
        }
    }

    fn run_cmd(program: &str, argv: &[~str], inputStr: Option<~str>, hasOutRedirect: bool) -> Option<~str> {
        if Shell::cmd_exists(program) {
            let mut options = run::ProcessOptions::new();
            options.in_fd = match inputStr {
                Some(_) => { None }
                None => { Some(0) }
            };

            options.out_fd = if hasOutRedirect {
                None
            } else {
                Some(1)
            };

            let mut process = run::Process::new(program, argv, options).unwrap();
            match inputStr {
                Some(string) => {
                    let buf = string.into_bytes();
                    process.input().write(buf);
                }
                None => {}
            }

            process.close_input();

            if hasOutRedirect {
                let processOutput = process.finish_with_output();
                return Some(str::from_utf8(processOutput.output).to_owned());
            } else {
                process.finish();
                return None;
            };
        } else {
            println!("{:s}: command not found", program);
            return None;
        }
    }

    fn cmd_exists(cmd_path: &str) -> bool {
        let ret = run::process_output("which", [cmd_path.to_owned()]);
        return ret.expect("exit code error.").status.success();
    }

    fn run_cd(cmd_line: &str) -> proc() {
        let params = cmd_line.to_owned();
        return proc() {
            let argv: ~[~str] = Shell::get_args(params);
            let pathOpt: Option<Path> = match argv.len() {
                1   =>  { os::homedir() }
                0   =>  { os::homedir() }
                _   =>  { Some(Path::new(argv[1])) }
            };

            match pathOpt {
                Some(path)   =>  {
                    if path.is_dir() {
                        os::change_dir(&path);
                    } else {
                        println!("Error: {:s} is not a directory", path.as_str().unwrap());
                    }
                }
                None        =>  {
                    println!("Error: Invalid path");
                }
            };
        };
    }

    fn run_history(cmd_line: &str, cmd_history: &[~str]) -> proc() {
        let history = cmd_history.to_owned();
        let params = cmd_line.to_owned();
        return proc() {
            let argv: ~[~str] = Shell::get_args(params);
            if (argv.len() > 1) {
                println!("Error: history does not take options (sadly).");
            }
            else {
                let mut i = 1;
                for entry in history.iter() {
                    println!("{:d} \t{:s}", i, entry.to_owned());
                    i = i + 1;
                }
            }
        };
    }

    fn split_on_last_pipe(cmd_line: &str) -> ~[~str] {
        return cmd_line.rsplitn('|', 1).filter_map(|x| Some(x.trim().to_owned())).to_owned_vec();
    }

    fn get_input_file_contents(cmd_line: &str) -> Option<~str> {
        return match cmd_line.splitn('<', 1).nth(1) {
            Some(inputFileCandidate) => {
                match inputFileCandidate.trim().splitn(' ', 1).nth(0) {
                    Some(inputFile) => {
                        match File::open(&Path::new(inputFile)) {
                            Some(mut file) => {
                                Some(file.read_to_str())
                            }
                            None => {
                                println!("File {:s} does not exist", inputFile);
                                None
                            }
                        }
                    }
                    None => { None }
                }
            }
            None => { None }
        };
    }

    fn get_output_file(cmd_line: &str) -> Option<~str> {
        return match cmd_line.splitn('>', 1).nth(1) {
            Some(inputFileCandidate) => {
                match inputFileCandidate.trim().splitn(' ', 1).nth(0) {
                    Some(inputFile) => { Some(inputFile.to_owned()) }
                    None => { None }
                }
            }
            None => { None }
        };
    }

    fn get_args_no_redirects(cmd_line: &str) -> ~[~str] {
        let mut activeRedirect = false;
        return cmd_line.split(' ').filter_map(|x| if x == ">" || x == "<" {
            activeRedirect = true;
            None
        } else if x == "" {
            None
        } else if activeRedirect {
            activeRedirect = false;
            None
        } else {
            Some(x.to_owned())
        }).to_owned_vec();
    }

    fn get_args(cmd_line: &str) -> ~[~str] {
        return cmd_line.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();
    }
}

fn get_cmdline_from_args() -> Option<~str> {
    /* Begin processing program arguments and initiate the parameters. */
    let args = os::args();

    let opts = ~[
        getopts::optopt("c")
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m }
        Err(f) => { fail!(f.to_err_msg()) }
    };

    if matches.opt_present("c") {
        let cmd_str = match matches.opt_str("c") {
                                                Some(cmd_str) => {cmd_str.to_owned()},
                                                None => {~""}
                                              };
        return Some(cmd_str);
    } else {
        return None;
    }
}

fn main() {
    let opt_cmd_line = get_cmdline_from_args();

    match opt_cmd_line {
        Some(cmd_line) => Shell::run_cmdline_single(cmd_line),
        None           => Shell::new("gash > ").run()
    }
}
<|endoftext|>