rust - Rust代码优化？

Question

我想写一个模拟 DNF 装备增幅的程序，通过多次样本执行得到平均每件增幅 10 装备需要增幅多少次。装备 +4 之前不会失败，+4 之后会失败且失败后还会掉级，具体如下图所示：

公会秘药和普雷宠物会额外增加每次增幅的成功率 1% 和 4%，所以一共分了三种情况。

我最开始用 js 写了一版：

const times = 1_000_000;
const data = [
  // 0 -> 1
  {
    rate: 1,
    failedTo: 0,
  },
  // 1 -> 2
  {
    rate: 1,
    failedTo: 1,
  },
  // 2 -> 3
  {
    rate: 1,
    failedTo: 2,
  },
  // 3 -> 4
  {
    rate: 1,
    failedTo: 3,
  },
  // 4 -> 5
  {
    rate: 0.8,
    failedTo: 3,
  },
  // 5 -> 6
  {
    rate: 0.7,
    failedTo: 4,
  },
  // 6 -> 7
  {
    rate: 0.6,
    failedTo: 5,
  },
  // 7 -> 8
  {
    rate: 0.7,
    failedTo: 4,
  },
  // 8 -> 9
  {
    rate: 0.6,
    failedTo: 5,
  },
  // 9 -> 10
  {
    rate: 0.5,
    failedTo: 6,
  },
];

function zengfuTest(title, ex = 0) {
  const log = {
    total: 0,
    0: 0,
    1: 0,
    2: 0,
    3: 0,
    4: 0,
    5: 0,
    6: 0,
    7: 0,
    8: 0,
    9: 0,
    _0: 0,
    _1: 0,
    _2: 0,
    _3: 0,
    _4: 0,
    _5: 0,
    _6: 0,
    _7: 0,
    _8: 0,
    _9: 0,
  };

  for (let i = 0; i < times; i++) {
    zengfuOne(0);
  }

  showLog(title);

  /**
   * @param {string} key
   */
  function logData(key) {
    log[key]++;
    log.total++;
  }

  /**
   * @param {number} rate
   */
  function zengfu(rate) {
    return Math.random() <= rate;
  }

  /**
   * @param {number} lv
   */
  function zengfuOne(lv) {
    if (lv === 10) return;

    const { rate, failedTo } = data[lv];
    const result = zengfu(rate + ex);

    if (result) {
      logData(lv);
      zengfuOne(lv + 1);
    } else {
      logData("_" + lv);
      zengfuOne(failedTo);
    }
  }

  function showLog(title) {
    console.log(`==== ${title} ====`);
    console.log(`测试数量：${times}`);
    console.log("平均每件红 10 所需增幅次数如下所示：");

    for (let i = 0; i <= 9; i++) {
      const success = (log[i] / times).toFixed(3);
      const fail = (log["_" + i] / times).toFixed(3);
      console.log(
        `增幅${i + 1}：成功 ${success} 次${
          log["_" + i] ? `，失败 ${fail} 次` : ""
        }`
      );
    }
    console.log(
      `\n平均每件红 10 需要增幅 ${(log.total / times).toFixed(3)} 次\n`
    );
  }
}

zengfuTest("红 10 增幅次数模拟");
zengfuTest("【公会秘药】红 10 增幅次数模拟", 0.01);
zengfuTest("【公会秘药】【普雷宠物】红 10 增幅次数模拟", 0.05);

后来想到我刚学了 rust，不如练练手，而且 rust 很快，于是又写了一版：

use rand::prelude::*;

struct Logger {
  title: String,
  total: i32,
  success: Vec<i32>,
  fail: Vec<i32>,
}

impl Logger {
  fn new(title: String) -> Self {
    Logger {
      title,
      total: 0,
      success: vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
      fail: vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    }
  }

  fn print(&self) {
    println!("==== {} ====", &self.title);
    println!("测试数量：{}", TEST_TIMES);
    println!("平均每件红 10 所需增幅次数如下所示：");

    let times: f64 = TEST_TIMES.try_into().unwrap();
    let total: f64 = self.total.try_into().unwrap();

    for i in 0..9 {
      let success: f64 = self.success[i].try_into().unwrap();
      let fail: f64 = self.fail[i].try_into().unwrap();

      println!(
        "增幅 {}：成功 {} 次，失败 {} 次",
        i + 1,
        success / times,
        fail / times
      );
    }

    println!("\n平均每件红 10 需要增幅 {} 次\n", total / times);
  }
}

#[derive(Debug)]
struct AmplifyData {
  rate: f64,
  fall: Option<i32>,
}

const TEST_TIMES: i32 = 1_000_000;

fn main() {
  test(String::from("红 10 增幅次数模拟"), 0.0);
  test(String::from("【公会秘药】红 10 增幅次数模拟"), 0.01);
  test(
    String::from("【公会秘药】【普雷宠物】红 10 增幅次数模拟"),
    0.05,
  );
}

fn test(title: String, ex: f64) {
  let mut logger = Logger::new(title);
  let data: Vec<AmplifyData> = vec![
    AmplifyData {
      rate: 1.0,
      fall: None,
    },
    AmplifyData {
      rate: 1.0,
      fall: None,
    },
    AmplifyData {
      rate: 1.0,
      fall: None,
    },
    AmplifyData {
      rate: 1.0,
      fall: None,
    },
    AmplifyData {
      rate: 0.8,
      fall: Some(3),
    },
    AmplifyData {
      rate: 0.7,
      fall: Some(4),
    },
    AmplifyData {
      rate: 0.6,
      fall: Some(5),
    },
    AmplifyData {
      rate: 0.7,
      fall: Some(4),
    },
    AmplifyData {
      rate: 0.6,
      fall: Some(5),
    },
    AmplifyData {
      rate: 0.5,
      fall: Some(6),
    },
  ];

  for _ in 0..TEST_TIMES {
    amplify_one(0, &data, ex, &mut logger);
  }

  logger.print();
}

fn amplify_once(rate: f64) -> bool {
  return thread_rng().gen::<f64>() <= rate;
}

fn amplify_one(lv: i32, data: &Vec<AmplifyData>, ex: f64, logger: &mut Logger) {
  if lv == 10 {
    return;
  }

  let index: usize = lv.try_into().unwrap();
  let item = &data[index];
  let AmplifyData { rate, fall, .. } = item;

  let true_rate = *rate + ex;

  let result = amplify_once(true_rate);

  logger.total += 1;
  let index: usize = lv.try_into().unwrap();

  if result {
    logger.success[index] += 1;
    amplify_one(lv + 1, data, ex, logger);
  } else {
    logger.fail[index] += 1;
    if let Some(next_lv) = fall {
      amplify_one(*next_lv, data, ex, logger);
    }
  }
}

然而实际上 rust 代码执行非常慢，js 那边三种模拟都走完了，这边还得等好几秒才出了第一个结果...不知道应该怎么优化一下我的代码？

Answer 1

优化了一下你看看：

use rand::prelude::*;
use std::convert::TryInto;

struct Logger {
    // ...
}

impl Logger {
    // ...
}

#[derive(Debug)]
struct AmplifyData {
    rate: f64,
    fall: Option<i32>,
}

const TEST_TIMES: i32 = 1_000_000;

fn main() {
    test(String::from("红 10 增幅次数模拟"), 0.0);
    test(String::from("【公会秘药】红 10 增幅次数模拟"), 0.01);
    test(
        String::from("【公会秘药】【普雷宠物】红 10 增幅次数模拟"),
        0.05,
    );
}

fn test(title: String, ex: f64) {
    let mut logger = Logger::new(title);
    let data: Vec<AmplifyData> = vec![
        // ...
    ];

    for _ in 0..TEST_TIMES {
        amplify_one(0, &data, ex, &mut logger);
    }

    logger.print();
}

fn amplify_once(rate: f64) -> bool {
    return thread_rng().gen::<f64>() <= rate;
}

fn amplify_one(lv: i32, data: &Vec<AmplifyData>, ex: f64, logger: &mut Logger) {
    let mut lv = lv;
    while lv != 10 {
        let index: usize = lv.try_into().unwrap();
        let item = &data[index];
        let AmplifyData { rate, fall, .. } = item;

        let true_rate = *rate + ex;

        let result = amplify_once(true_rate);

        logger.total += 1;
        let index: usize = lv.try_into().unwrap();

        if result {
            logger.success[index] += 1;
            lv += 1;
        } else {
            logger.fail[index] += 1;
            if let Some(next_lv) = fall {
                lv = *next_lv;
            }
        }
    }
}

用 rayon库实现并行计算再优化：

use rand::prelude::*;
use rayon::prelude::*;
use std::convert::TryInto;

struct Logger {
    // ...
}

impl Logger {
    // ...
}

#[derive(Debug, Clone, Copy)]
struct AmplifyData {
    rate: f64,
    fall: Option<i32>,
}

const TEST_TIMES: i32 = 1_000_000;

fn main() {
    test(String::from("红 10 增幅次数模拟"), 0.0);
    test(String::from("【公会秘药】红 10 增幅次数模拟"), 0.01);
    test(
        String::from("【公会秘药】【普雷宠物】红 10 增幅次数模拟"),
        0.05,
    );
}

fn test(title: String, ex: f64) {
    let mut logger = Logger::new(title);
    let data: Vec<AmplifyData> = vec![
        // ...
    ];

    let num_threads = rayon::current_num_threads();
    let chunk_size = TEST_TIMES / num_threads as i32;

    (0..num_threads)
        .into_par_iter()
        .for_each(|_| {
            let mut thread_local_logger = Logger::new("".to_string());
            for _ in 0..chunk_size {
                amplify_one(0, &data, ex, &mut thread_local_logger);
            }
            logger.merge(&thread_local_logger);
        });

    logger.print();
}

impl Logger {
    fn merge(&mut self, other: &Self) {
        self.total += other.total;
        for i in 0..10 {
            self.success[i] += other.success[i];
            self.fail[i] += other.fail[i];
        }
    }
}

fn amplify_once(rate: f64) -> bool {
    return thread_rng().gen::<f64>() <= rate;
}

fn amplify_one(lv: i32, data: &Vec<AmplifyData>, ex: f64, logger: &mut Logger) {
    let mut lv = lv;
    while lv != 10 {
        // ...
    }
}