<commit_before>use compression::{NUM_DOCS_PER_BLOCK, BlockDecoder, VIntDecoder};
use DocId;
use postings::{Postings, FreqHandler, DocSet, HasLen, SkipResult};
use std::num::Wrapping;
use fastfield::DeleteBitSet;

const EMPTY_DATA: [u8; 0] = [0u8; 0];

/// `SegmentPostings` represents the inverted list or postings associated to
/// a term in a `Segment`.
///
/// As we iterate through the `SegmentPostings`, the frequencies are optionally decoded.
/// Positions on the other hand, are optionally entirely decoded upfront.
pub struct SegmentPostings<'a> {
    len: usize,
    block_len: usize,
    doc_offset: u32,
    block_decoder: BlockDecoder,
    freq_handler: FreqHandler,
    remaining_data: &'a [u8],
    cur: Wrapping<usize>,
    delete_bitset: DeleteBitSet,
}

impl<'a> SegmentPostings<'a> {
    fn load_next_block(&mut self) {
        let num_remaining_docs = self.len - self.cur.0;
        if num_remaining_docs >= NUM_DOCS_PER_BLOCK {
            self.remaining_data =
                self.block_decoder
                    .uncompress_block_sorted(self.remaining_data, self.doc_offset);
            self.remaining_data = self.freq_handler.read_freq_block(self.remaining_data);
            self.doc_offset = self.block_decoder.output(NUM_DOCS_PER_BLOCK - 1);
            self.block_len = NUM_DOCS_PER_BLOCK;
        } else {
            self.remaining_data =
                self.block_decoder
                    .uncompress_vint_sorted(self.remaining_data,
                                            self.doc_offset,
                                            num_remaining_docs);
            self.freq_handler
                .read_freq_vint(self.remaining_data, num_remaining_docs);
            self.block_len = num_remaining_docs;
        }
    }

    /// Reads a Segment postings from an &[u8]
    ///
    /// * `len` - number of document in the posting lists.
    /// * `data` - data array. The complete data is not necessarily used.
    /// * `freq_handler` - the freq handler is in charge of decoding
    ///   frequencies and/or positions
    pub fn from_data(len: u32,
                     data: &'a [u8],
                     delete_bitset: &'a DeleteBitSet,
                     freq_handler: FreqHandler)
                     -> SegmentPostings<'a> {
        SegmentPostings {
            len: len as usize,
            block_len: len as usize,
            doc_offset: 0,
            block_decoder: BlockDecoder::new(),
            freq_handler: freq_handler,
            remaining_data: data,
            cur: Wrapping(usize::max_value()),
            delete_bitset: delete_bitset.clone(),
        }
    }

    /// Returns an empty segment postings object
    pub fn empty() -> SegmentPostings<'static> {
        SegmentPostings {
            len: 0,
            block_len: 0,
            doc_offset: 0,
            block_decoder: BlockDecoder::new(),
            freq_handler: FreqHandler::new_without_freq(),
            remaining_data: &EMPTY_DATA,
            delete_bitset: DeleteBitSet::empty(),
            cur: Wrapping(usize::max_value()),
        }
    }

    /// Index within a block is used as an address when
    /// interacting with the `FreqHandler`
    fn index_within_block(&self) -> usize {
        self.cur.0 % NUM_DOCS_PER_BLOCK
    }

    /// Sets the current position to a location relative
    /// to the current block
    #[inline]
    fn set_within_block(&mut self, inner_pos: usize) {
        self.cur = Wrapping(self.cur.0 & !(NUM_DOCS_PER_BLOCK - 1)) + Wrapping(inner_pos)
    }
}

impl<'a> DocSet for SegmentPostings<'a> {
    // goes to the next element.
    // next needs to be called a first time to point to the correct element.
    #[inline]
    fn advance(&mut self) -> bool {
        loop {
            self.cur += Wrapping(1);
            if self.cur.0 >= self.len {
                return false;
            }
            if self.index_within_block() == 0 {
                self.load_next_block();
            }
            if !self.delete_bitset.is_deleted(self.doc()) {
                return true;
            }
        }
    }

    fn skip_next(&mut self, target: DocId) -> SkipResult {
        if !self.advance() {
            return SkipResult::End;
        }

        let mut pos = self.index_within_block();
        // skip blocks until one that might contain the target
        loop {
            // check if we need to go to the next block
            if target > self.block_decoder.output(self.block_len - 1) {
                self.cur += Wrapping(self.block_len - pos);
                self.load_next_block();
                pos = 0;

                // there was no more data
                if self.cur.0 == self.len {
                    return SkipResult::End;
                }
            } else if target < self.block_decoder.output(pos) {
                // We've overpassed the target after the first `advance` call
                // or we're at the beginning of a block.
                // Either way, we're on the first `DocId` greater than `target`
                return SkipResult::OverStep;
            } else {
                break;
            }
        }

        debug_assert!(target >= self.block_decoder.output(pos));
        debug_assert!(target <= self.block_decoder.output(self.block_len - 1));

        // we're in the right block, do a binary search
        let mut start = pos;
        let mut count = self.block_len - start;
        while count > 0 {
            let step = count / 2;
            let mid = start + step;
            let doc = self.block_decoder.output(mid);
            if doc < target {
                start = mid + 1;
                count -= step + 1;
            } else {
                count = step;
            }
        }

        // `doc` is now >= `target`
        let doc = self.block_decoder.output(start);
        self.set_within_block(start);

        if !self.delete_bitset.is_deleted(doc) {
            if doc == target {
                return SkipResult::Reached;
            } else {
                return SkipResult::OverStep;
            }
        }

        if self.advance() {
            return SkipResult::OverStep;
        }
        return SkipResult::End;
    }

    #[inline]
    fn doc(&self) -> DocId {
        self.block_decoder.output(self.index_within_block())
    }
}

impl<'a> HasLen for SegmentPostings<'a> {
    fn len(&self) -> usize {
        self.len
    }
}

impl<'a> Postings for SegmentPostings<'a> {
    fn term_freq(&self) -> u32 {
        self.freq_handler.freq(self.index_within_block())
    }

    fn positions(&self) -> &[u32] {
        self.freq_handler.positions(self.index_within_block())
    }
}
<commit_msg>Add comment about block_len<commit_after>use compression::{NUM_DOCS_PER_BLOCK, BlockDecoder, VIntDecoder};
use DocId;
use postings::{Postings, FreqHandler, DocSet, HasLen, SkipResult};
use std::num::Wrapping;
use fastfield::DeleteBitSet;

const EMPTY_DATA: [u8; 0] = [0u8; 0];

/// `SegmentPostings` represents the inverted list or postings associated to
/// a term in a `Segment`.
///
/// As we iterate through the `SegmentPostings`, the frequencies are optionally decoded.
/// Positions on the other hand, are optionally entirely decoded upfront.
pub struct SegmentPostings<'a> {
    len: usize,
    // Removing this makes the code slower
    // See https://github.com/tantivy-search/tantivy/issues/89
    block_len: usize,
    doc_offset: u32,
    block_decoder: BlockDecoder,
    freq_handler: FreqHandler,
    remaining_data: &'a [u8],
    cur: Wrapping<usize>,
    delete_bitset: DeleteBitSet,
}

impl<'a> SegmentPostings<'a> {
    fn load_next_block(&mut self) {
        let num_remaining_docs = self.len - self.cur.0;
        if num_remaining_docs >= NUM_DOCS_PER_BLOCK {
            self.remaining_data =
                self.block_decoder
                    .uncompress_block_sorted(self.remaining_data, self.doc_offset);
            self.remaining_data = self.freq_handler.read_freq_block(self.remaining_data);
            self.doc_offset = self.block_decoder.output(NUM_DOCS_PER_BLOCK - 1);
            self.block_len = NUM_DOCS_PER_BLOCK;
        } else {
            self.remaining_data =
                self.block_decoder
                    .uncompress_vint_sorted(self.remaining_data,
                                            self.doc_offset,
                                            num_remaining_docs);
            self.freq_handler
                .read_freq_vint(self.remaining_data, num_remaining_docs);
            self.block_len = num_remaining_docs;
        }
    }

    /// Reads a Segment postings from an &[u8]
    ///
    /// * `len` - number of document in the posting lists.
    /// * `data` - data array. The complete data is not necessarily used.
    /// * `freq_handler` - the freq handler is in charge of decoding
    ///   frequencies and/or positions
    pub fn from_data(len: u32,
                     data: &'a [u8],
                     delete_bitset: &'a DeleteBitSet,
                     freq_handler: FreqHandler)
                     -> SegmentPostings<'a> {
        SegmentPostings {
            len: len as usize,
            block_len: len as usize,
            doc_offset: 0,
            block_decoder: BlockDecoder::new(),
            freq_handler: freq_handler,
            remaining_data: data,
            cur: Wrapping(usize::max_value()),
            delete_bitset: delete_bitset.clone(),
        }
    }

    /// Returns an empty segment postings object
    pub fn empty() -> SegmentPostings<'static> {
        SegmentPostings {
            len: 0,
            block_len: 0,
            doc_offset: 0,
            block_decoder: BlockDecoder::new(),
            freq_handler: FreqHandler::new_without_freq(),
            remaining_data: &EMPTY_DATA,
            delete_bitset: DeleteBitSet::empty(),
            cur: Wrapping(usize::max_value()),
        }
    }

    /// Index within a block is used as an address when
    /// interacting with the `FreqHandler`
    fn index_within_block(&self) -> usize {
        self.cur.0 % NUM_DOCS_PER_BLOCK
    }

    /// Sets the current position to a location relative
    /// to the current block
    #[inline]
    fn set_within_block(&mut self, inner_pos: usize) {
        self.cur = Wrapping(self.cur.0 & !(NUM_DOCS_PER_BLOCK - 1)) + Wrapping(inner_pos)
    }
}

impl<'a> DocSet for SegmentPostings<'a> {
    // goes to the next element.
    // next needs to be called a first time to point to the correct element.
    #[inline]
    fn advance(&mut self) -> bool {
        loop {
            self.cur += Wrapping(1);
            if self.cur.0 >= self.len {
                return false;
            }
            if self.index_within_block() == 0 {
                self.load_next_block();
            }
            if !self.delete_bitset.is_deleted(self.doc()) {
                return true;
            }
        }
    }

    fn skip_next(&mut self, target: DocId) -> SkipResult {
        if !self.advance() {
            return SkipResult::End;
        }

        let mut pos = self.index_within_block();
        // skip blocks until one that might contain the target
        loop {
            // check if we need to go to the next block
            if target > self.block_decoder.output(self.block_len - 1) {
                self.cur += Wrapping(self.block_len - pos);
                self.load_next_block();
                pos = 0;

                // there was no more data
                if self.cur.0 == self.len {
                    return SkipResult::End;
                }
            } else if target < self.block_decoder.output(pos) {
                // We've overpassed the target after the first `advance` call
                // or we're at the beginning of a block.
                // Either way, we're on the first `DocId` greater than `target`
                return SkipResult::OverStep;
            } else {
                break;
            }
        }

        debug_assert!(target >= self.block_decoder.output(pos));
        debug_assert!(target <= self.block_decoder.output(self.block_len - 1));

        // we're in the right block, do a binary search
        let mut start = pos;
        let mut count = self.block_len - start;
        while count > 0 {
            let step = count / 2;
            let mid = start + step;
            let doc = self.block_decoder.output(mid);
            if doc < target {
                start = mid + 1;
                count -= step + 1;
            } else {
                count = step;
            }
        }

        // `doc` is now >= `target`
        let doc = self.block_decoder.output(start);
        self.set_within_block(start);

        if !self.delete_bitset.is_deleted(doc) {
            if doc == target {
                return SkipResult::Reached;
            } else {
                return SkipResult::OverStep;
            }
        }

        if self.advance() {
            return SkipResult::OverStep;
        }
        return SkipResult::End;
    }

    #[inline]
    fn doc(&self) -> DocId {
        self.block_decoder.output(self.index_within_block())
    }
}

impl<'a> HasLen for SegmentPostings<'a> {
    fn len(&self) -> usize {
        self.len
    }
}

impl<'a> Postings for SegmentPostings<'a> {
    fn term_freq(&self) -> u32 {
        self.freq_handler.freq(self.index_within_block())
    }

    fn positions(&self) -> &[u32] {
        self.freq_handler.positions(self.index_within_block())
    }
}
<|endoftext|>


<commit_before>use std::f32::consts::PI;
use std::sync::{Arc, Mutex,
                mpsc::{Sender, Receiver, TryRecvError},
                atomic::{AtomicUsize, Ordering},
};

use cgmath::{Point3, Vector3, Vector4, prelude::*};

use glium::Rect;

use rand::{self, prelude::*};

use crate::bvh::BVHNode;
use crate::camera::Camera;
use crate::color::Color;
use crate::material::Material;
use crate::pt_renderer::{Intersect, Ray, RenderCoordinator};
use crate::scene::Scene;
use crate::triangle::Hit;

const EPSILON: f32 = 1e-5;
const _EB: f32 = 5.0; // Desired expectation value of bounces
// The matching survival probability from negative binomial distribution
const RR_PROB: f32 = _EB / (_EB + 1.0);

pub struct RenderWorker {
    scene: Arc<Scene>,
    camera: Camera,
    coordinator: Arc<Mutex<RenderCoordinator>>,
    message_rx: Receiver<()>,
    result_tx: Sender<(Rect, Vec<f32>)>,
    ray_count: Arc<AtomicUsize>,
}

impl RenderWorker {
    pub fn new(scene: Arc<Scene>, camera: Camera, coordinator: Arc<Mutex<RenderCoordinator>>,
               message_rx: Receiver<()>, result_tx: Sender<(Rect, Vec<f32>)>,
               ray_count: Arc<AtomicUsize>) -> RenderWorker {
        RenderWorker {
            scene, camera, coordinator,
            message_rx, result_tx,
            ray_count,
        }
    }

    pub fn run(&self) {
        let (width, height) = {
            let coordinator = self.coordinator.lock().unwrap();
            (coordinator.width, coordinator.height)
        };
        let clip_to_world = self.camera.get_world_to_clip().invert().unwrap();
        let mut node_stack: Vec<(&BVHNode, f32)> = Vec::new();
        loop {
            match self.message_rx.try_recv() {
                Err(TryRecvError::Empty) => (),
                Ok(_) => return,
                Err(TryRecvError::Disconnected) => {
                    println!("Threads were not properly stopped before disconnecting channel!");
                    return;
                }
            }
            let block = {
                let mut coordinator = self.coordinator.lock().unwrap();
                coordinator.next_block()
            };
            if let Some(rect) = block {
                let mut block = vec![0.0f32; (3 * rect.width * rect.height) as usize];
                for h in 0..rect.height {
                    for w in 0..rect.width {
                        let samples_per_dir = 2u32;
                        let mut c = Color::black();
                        for j in 0..samples_per_dir {
                            for i in 0..samples_per_dir {
                                let dx = (i as f32 + rand::random::<f32>()) / samples_per_dir as f32;
                                let dy = (j as f32 + rand::random::<f32>()) / samples_per_dir as f32;
                                let clip_x = 2.0 * ((rect.left + w) as f32 + dx) / width as f32 - 1.0;
                                let clip_y = 2.0 * ((rect.bottom + h) as f32 + dy) / height as f32 - 1.0;
                                let clip_p = Vector4::new(clip_x, clip_y, 1.0, 1.0);
                                let world_p = clip_to_world * clip_p;
                                let dir = ((world_p / world_p.w).truncate() - self.camera.pos.to_vec())
                                    .normalize();
                                let ray = Ray::new(self.camera.pos, dir, 10.0 * self.camera.scale);
                                c += self.trace_ray(&ray, &mut node_stack, 0);
                            }
                        }
                        c /= samples_per_dir.pow(2) as f32;
                        let pixel_i = 3 * (h * rect.width + w) as usize;
                        block[pixel_i]     = c.r;
                        block[pixel_i + 1] = c.g;
                        block[pixel_i + 2] = c.b;
                    }
                }
                self.result_tx.send((rect, block)).expect("Receiver closed!");
            } else {
                return;
            }
        }
    }

    fn trace_ray(&'a self, ray: &Ray, node_stack: &mut Vec<(&'a BVHNode, f32)>, bounce: u32) -> Color {
        let mut c = Color::black();
        if let Some(hit) = self.find_hit(&ray, node_stack) {
            let material = &self.scene.materials[hit.tri.material_i];
            let mut normal = hit.normal();
            // Flip the normal if its pointing to the opposite side from the hit
            if normal.dot(ray.dir) > 0.0 {
                normal *= -1.0;
            }
            if bounce == 0 {
                if let Some(emissive) = material.emissive {
                    c += normal.dot(-ray.dir).max(0.0) * emissive;
                }
            }
            let (emissive, light_pos, light_normal, light_pdf) = self.sample_light();
            let bump_pos = hit.pos() + EPSILON * normal;
            let hit_to_light = light_pos - bump_pos;
            let light_dir = hit_to_light.normalize();
            let shadow_ray = Ray::new(bump_pos, light_dir, hit_to_light.magnitude() - EPSILON);
            if shadow_ray.length > 0.1 && self.find_hit(&shadow_ray, node_stack).is_none() {
                let cos_l = light_normal.dot(-light_dir).max(0.0);
                let cos_t = normal.dot(light_dir).max(0.0);
                c += emissive * self.brdf(&hit, &ray, &shadow_ray, material)
                    * cos_l * cos_t / (hit_to_light.magnitude2() * light_pdf);
            }
            let rr: f32 = rand::random();
            if rr < RR_PROB {
                let (new_dir, mut pdf) = self.sample_dir(normal);
                pdf *= RR_PROB;
                let new_ray = Ray::new(bump_pos, new_dir, 10.0 * self.camera.scale);
                c += normal.dot(new_dir) * self.brdf(&hit, &ray, &new_ray, material)
                    * self.trace_ray(&new_ray, node_stack, bounce + 1) / pdf;
            }
        }
        c
    }

    fn brdf(&self, hit: &Hit, _in_ray: &Ray, _out_ray: &Ray, material: &Material) -> Color {
        let tex_coords = hit.tex_coords();
        material.diffuse(tex_coords) / PI
    }

    fn sample_dir(&self, normal: Vector3<f32>) -> (Vector3<f32>, f32) {
        let dir = 2.0 * PI * rand::random::<f32>();
        let length: f32 = rand::random();
        let x = length * dir.cos();
        let y = length * dir.sin();
        let z = (1.0 - length.powi(2)).sqrt();
        let nx = if normal.x.abs() > normal.y.abs() {
            Vector3::new(normal.z, 0.0, -normal.x).normalize()
        } else {
            Vector3::new(0.0, -normal.z, normal.y).normalize()
        };
        let ny = normal.cross(nx);
        (x * nx + y * ny + z * normal, z / PI)
    }

    pub fn sample_light(&self) -> (Color, Point3<f32>, Vector3<f32>, f32) {
        if self.scene.lights.is_empty() {
            let light_intensity = 10.0 * self.camera.scale;
            (light_intensity * Color::white(), self.camera.pos, self.camera.dir, 1.0)
        } else {
            let i = rand::thread_rng().gen_range(0, self.scene.lights.len());
            let light = &self.scene.lights[i];
            let pdf = 1.0 / (self.scene.lights.len() as f32 * light.area());
            let (point, normal) = light.random_point();
            let light_material = &self.scene.materials[light.material_i];
            let emissive = light_material.emissive.expect("Light wasn't emissive");
            (emissive, point, normal, pdf)
        }
    }

    fn find_hit(&'a self, ray: &Ray, node_stack: &mut Vec<(&'a BVHNode, f32)>) -> Option<Hit> {
        self.ray_count.fetch_add(1, Ordering::Relaxed);
        let bvh = &self.scene.bvh;
        node_stack.push((bvh.root(), 0.0f32));
        let mut closest_hit: Option<Hit> = None;
        while let Some((node, t)) = node_stack.pop() {
            // We've already found closer hit
            if closest_hit.as_ref().map_or(false, |closest| closest.t <= t) { continue }
            if node.is_leaf() {
                for tri in &self.scene.triangles[node.start_i..node.end_i] {
                    if let Some(hit) = tri.intersect(&ray) {
                        if let Some(closest) = closest_hit.take() {
                            if hit.t < closest.t {
                                closest_hit = Some(hit);
                            } else {
                                closest_hit = Some(closest);
                            }
                        } else {
                            closest_hit = Some(hit);
                        }
                    }
                }
            } else {
                let (left, right) = bvh.get_children(node)
                    .expect("Non leaf node had no child nodes!");
                // TODO: benchmark this properly on the larger scenes
                if false { // not sorted
                    if let Some(t_left) = left.intersect(&ray) {
                        node_stack.push((left, t_left));
                    }
                    if let Some(t_right) = right.intersect(&ray) {
                        node_stack.push((right, t_right));
                    }
                } else { // sort nodes
                    let left_intersect = left.intersect(&ray);
                    let right_intersect = right.intersect(&ray);
                    if let Some(t_left) = left_intersect {
                        if let Some(t_right) = right_intersect {
                            // Put the closer hit on top
                            if t_left >= t_right {
                                node_stack.push((left, t_left));
                                node_stack.push((right, t_right));
                            } else {
                                node_stack.push((right, t_right));
                                node_stack.push((left, t_left));
                            }
                        } else {
                            node_stack.push((left, t_left));
                        }
                    } else if let Some(t_right) = right_intersect {
                        node_stack.push((right, t_right));
                    }
                }
            }
        }
        closest_hit
    }
}
<commit_msg>Allow short shadow rays<commit_after>use std::f32::consts::PI;
use std::sync::{Arc, Mutex,
                mpsc::{Sender, Receiver, TryRecvError},
                atomic::{AtomicUsize, Ordering},
};

use cgmath::{Point3, Vector3, Vector4, prelude::*};

use glium::Rect;

use rand::{self, prelude::*};

use crate::bvh::BVHNode;
use crate::camera::Camera;
use crate::color::Color;
use crate::material::Material;
use crate::pt_renderer::{Intersect, Ray, RenderCoordinator};
use crate::scene::Scene;
use crate::triangle::Hit;

const EPSILON: f32 = 1e-5;
const _EB: f32 = 5.0; // Desired expectation value of bounces
// The matching survival probability from negative binomial distribution
const RR_PROB: f32 = _EB / (_EB + 1.0);

pub struct RenderWorker {
    scene: Arc<Scene>,
    camera: Camera,
    coordinator: Arc<Mutex<RenderCoordinator>>,
    message_rx: Receiver<()>,
    result_tx: Sender<(Rect, Vec<f32>)>,
    ray_count: Arc<AtomicUsize>,
}

impl RenderWorker {
    pub fn new(scene: Arc<Scene>, camera: Camera, coordinator: Arc<Mutex<RenderCoordinator>>,
               message_rx: Receiver<()>, result_tx: Sender<(Rect, Vec<f32>)>,
               ray_count: Arc<AtomicUsize>) -> RenderWorker {
        RenderWorker {
            scene, camera, coordinator,
            message_rx, result_tx,
            ray_count,
        }
    }

    pub fn run(&self) {
        let (width, height) = {
            let coordinator = self.coordinator.lock().unwrap();
            (coordinator.width, coordinator.height)
        };
        let clip_to_world = self.camera.get_world_to_clip().invert().unwrap();
        let mut node_stack: Vec<(&BVHNode, f32)> = Vec::new();
        loop {
            match self.message_rx.try_recv() {
                Err(TryRecvError::Empty) => (),
                Ok(_) => return,
                Err(TryRecvError::Disconnected) => {
                    println!("Threads were not properly stopped before disconnecting channel!");
                    return;
                }
            }
            let block = {
                let mut coordinator = self.coordinator.lock().unwrap();
                coordinator.next_block()
            };
            if let Some(rect) = block {
                let mut block = vec![0.0f32; (3 * rect.width * rect.height) as usize];
                for h in 0..rect.height {
                    for w in 0..rect.width {
                        let samples_per_dir = 2u32;
                        let mut c = Color::black();
                        for j in 0..samples_per_dir {
                            for i in 0..samples_per_dir {
                                let dx = (i as f32 + rand::random::<f32>()) / samples_per_dir as f32;
                                let dy = (j as f32 + rand::random::<f32>()) / samples_per_dir as f32;
                                let clip_x = 2.0 * ((rect.left + w) as f32 + dx) / width as f32 - 1.0;
                                let clip_y = 2.0 * ((rect.bottom + h) as f32 + dy) / height as f32 - 1.0;
                                let clip_p = Vector4::new(clip_x, clip_y, 1.0, 1.0);
                                let world_p = clip_to_world * clip_p;
                                let dir = ((world_p / world_p.w).truncate() - self.camera.pos.to_vec())
                                    .normalize();
                                let ray = Ray::new(self.camera.pos, dir, 10.0 * self.camera.scale);
                                c += self.trace_ray(&ray, &mut node_stack, 0);
                            }
                        }
                        c /= samples_per_dir.pow(2) as f32;
                        let pixel_i = 3 * (h * rect.width + w) as usize;
                        block[pixel_i]     = c.r;
                        block[pixel_i + 1] = c.g;
                        block[pixel_i + 2] = c.b;
                    }
                }
                self.result_tx.send((rect, block)).expect("Receiver closed!");
            } else {
                return;
            }
        }
    }

    fn trace_ray(&'a self, ray: &Ray, node_stack: &mut Vec<(&'a BVHNode, f32)>, bounce: u32) -> Color {
        let mut c = Color::black();
        if let Some(hit) = self.find_hit(&ray, node_stack) {
            let material = &self.scene.materials[hit.tri.material_i];
            let mut normal = hit.normal();
            // Flip the normal if its pointing to the opposite side from the hit
            if normal.dot(ray.dir) > 0.0 {
                normal *= -1.0;
            }
            if bounce == 0 {
                if let Some(emissive) = material.emissive {
                    c += normal.dot(-ray.dir).max(0.0) * emissive;
                }
            }
            let (emissive, light_pos, light_normal, light_pdf) = self.sample_light();
            let bump_pos = hit.pos() + EPSILON * normal;
            let hit_to_light = light_pos - bump_pos;
            let light_dir = hit_to_light.normalize();
            let shadow_ray = Ray::new(bump_pos, light_dir, hit_to_light.magnitude() - EPSILON);
            if self.find_hit(&shadow_ray, node_stack).is_none() {
                let cos_l = light_normal.dot(-light_dir).max(0.0);
                let cos_t = normal.dot(light_dir).max(0.0);
                c += emissive * self.brdf(&hit, &ray, &shadow_ray, material)
                    * cos_l * cos_t / (hit_to_light.magnitude2() * light_pdf);
            }
            let rr: f32 = rand::random();
            if rr < RR_PROB {
                let (new_dir, mut pdf) = self.sample_dir(normal);
                pdf *= RR_PROB;
                let new_ray = Ray::new(bump_pos, new_dir, 10.0 * self.camera.scale);
                c += normal.dot(new_dir) * self.brdf(&hit, &ray, &new_ray, material)
                    * self.trace_ray(&new_ray, node_stack, bounce + 1) / pdf;
            }
        }
        c
    }

    fn brdf(&self, hit: &Hit, _in_ray: &Ray, _out_ray: &Ray, material: &Material) -> Color {
        let tex_coords = hit.tex_coords();
        material.diffuse(tex_coords) / PI
    }

    fn sample_dir(&self, normal: Vector3<f32>) -> (Vector3<f32>, f32) {
        let dir = 2.0 * PI * rand::random::<f32>();
        let length: f32 = rand::random();
        let x = length * dir.cos();
        let y = length * dir.sin();
        let z = (1.0 - length.powi(2)).sqrt();
        let nx = if normal.x.abs() > normal.y.abs() {
            Vector3::new(normal.z, 0.0, -normal.x).normalize()
        } else {
            Vector3::new(0.0, -normal.z, normal.y).normalize()
        };
        let ny = normal.cross(nx);
        (x * nx + y * ny + z * normal, z / PI)
    }

    pub fn sample_light(&self) -> (Color, Point3<f32>, Vector3<f32>, f32) {
        if self.scene.lights.is_empty() {
            let light_intensity = 10.0 * self.camera.scale;
            (light_intensity * Color::white(), self.camera.pos, self.camera.dir, 1.0)
        } else {
            let i = rand::thread_rng().gen_range(0, self.scene.lights.len());
            let light = &self.scene.lights[i];
            let pdf = 1.0 / (self.scene.lights.len() as f32 * light.area());
            let (point, normal) = light.random_point();
            let light_material = &self.scene.materials[light.material_i];
            let emissive = light_material.emissive.expect("Light wasn't emissive");
            (emissive, point, normal, pdf)
        }
    }

    fn find_hit(&'a self, ray: &Ray, node_stack: &mut Vec<(&'a BVHNode, f32)>) -> Option<Hit> {
        self.ray_count.fetch_add(1, Ordering::Relaxed);
        let bvh = &self.scene.bvh;
        node_stack.push((bvh.root(), 0.0f32));
        let mut closest_hit: Option<Hit> = None;
        while let Some((node, t)) = node_stack.pop() {
            // We've already found closer hit
            if closest_hit.as_ref().map_or(false, |closest| closest.t <= t) { continue }
            if node.is_leaf() {
                for tri in &self.scene.triangles[node.start_i..node.end_i] {
                    if let Some(hit) = tri.intersect(&ray) {
                        if let Some(closest) = closest_hit.take() {
                            if hit.t < closest.t {
                                closest_hit = Some(hit);
                            } else {
                                closest_hit = Some(closest);
                            }
                        } else {
                            closest_hit = Some(hit);
                        }
                    }
                }
            } else {
                let (left, right) = bvh.get_children(node)
                    .expect("Non leaf node had no child nodes!");
                // TODO: benchmark this properly on the larger scenes
                if false { // not sorted
                    if let Some(t_left) = left.intersect(&ray) {
                        node_stack.push((left, t_left));
                    }
                    if let Some(t_right) = right.intersect(&ray) {
                        node_stack.push((right, t_right));
                    }
                } else { // sort nodes
                    let left_intersect = left.intersect(&ray);
                    let right_intersect = right.intersect(&ray);
                    if let Some(t_left) = left_intersect {
                        if let Some(t_right) = right_intersect {
                            // Put the closer hit on top
                            if t_left >= t_right {
                                node_stack.push((left, t_left));
                                node_stack.push((right, t_right));
                            } else {
                                node_stack.push((right, t_right));
                                node_stack.push((left, t_left));
                            }
                        } else {
                            node_stack.push((left, t_left));
                        }
                    } else if let Some(t_right) = right_intersect {
                        node_stack.push((right, t_right));
                    }
                }
            }
        }
        closest_hit
    }
}
<|endoftext|>


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

use matches::assert_matches;
use ruma_events::{
    custom::RedactedCustomEventContent,
    room::{
        aliases::RedactedAliasesEventContent,
        create::RedactedCreateEventContent,
        message::RedactedMessageEventContent,
        redaction::{RedactionEvent, RedactionEventContent},
    },
    AnyRedactedMessageEvent, AnyRedactedSyncMessageEvent, AnyRedactedSyncStateEvent, AnyRoomEvent,
    AnySyncRoomEvent, EventJson, RedactedMessageEvent, RedactedSyncMessageEvent,
    RedactedSyncStateEvent, UnsignedData,
};
use ruma_identifiers::{EventId, RoomId, UserId};
use serde_json::{from_value as from_json_value, json, to_value as to_json_value};

fn is_zst<T>(_: &T) -> bool {
    std::mem::size_of::<T>() == 0
}

#[test]
fn redacted_message_event_serialize() {
    let redacted = RedactedSyncMessageEvent {
        content: RedactedMessageEventContent,
        event_id: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        origin_server_ts: UNIX_EPOCH + Duration::from_millis(1),
        sender: UserId::try_from("@carl:example.com").unwrap(),
        unsigned: UnsignedData::default(),
    };

    let expected = json!({
      "event_id": "$h29iv0s8:example.com",
      "origin_server_ts": 1,
      "sender": "@carl:example.com",
      "type": "m.room.message"
    });

    let actual = to_json_value(&redacted).unwrap();
    assert_eq!(actual, expected);
}

#[test]
fn redacted_aliases_event_serialize() {
    let redacted = RedactedSyncStateEvent {
        content: RedactedAliasesEventContent { aliases: None },
        event_id: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        state_key: "".into(),
        origin_server_ts: UNIX_EPOCH + Duration::from_millis(1),
        sender: UserId::try_from("@carl:example.com").unwrap(),
        unsigned: UnsignedData::default(),
    };

    let expected = json!({
      "event_id": "$h29iv0s8:example.com",
      "state_key": "",
      "origin_server_ts": 1,
      "sender": "@carl:example.com",
      "type": "m.room.aliases"
    });

    let actual = to_json_value(&redacted).unwrap();
    assert_eq!(actual, expected);
}

#[test]
fn redacted_deserialize_any_room() {
    let mut unsigned = UnsignedData::default();
    // The presence of `redacted_because` triggers the event enum (AnyRoomEvent in this case)
    // to return early with `RedactedContent` instead of failing to deserialize according
    // to the event type string.
    unsigned.redacted_because = Some(EventJson::from(RedactionEvent {
        content: RedactionEventContent { reason: Some("redacted because".into()) },
        redacts: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        event_id: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        origin_server_ts: UNIX_EPOCH + Duration::from_millis(1),
        room_id: RoomId::try_from("!roomid:room.com").unwrap(),
        sender: UserId::try_from("@carl:example.com").unwrap(),
        unsigned: UnsignedData::default(),
    }));

    let redacted = json!({
      "event_id": "$h29iv0s8:example.com",
      "room_id": "!roomid:room.com",
      "origin_server_ts": 1,
      "sender": "@carl:example.com",
      "unsigned": unsigned,
      "type": "m.room.message"
    });

    let actual = to_json_value(&redacted).unwrap();

    assert_matches!(
        from_json_value::<EventJson<AnyRoomEvent>>(actual)
            .unwrap()
            .deserialize()
            .unwrap(),
        AnyRoomEvent::RedactedMessage(AnyRedactedMessageEvent::RoomMessage(RedactedMessageEvent {
            event_id, room_id, content, ..
        })) if event_id == EventId::try_from("$h29iv0s8:example.com").unwrap()
            && room_id == RoomId::try_from("!roomid:room.com").unwrap()
            && is_zst(&content)
    )
}

#[test]
fn redacted_deserialize_any_room_sync() {
    let mut unsigned = UnsignedData::default();
    // The presence of `redacted_because` triggers the event enum (AnySyncRoomEvent in this case)
    // to return early with `RedactedContent` instead of failing to deserialize according
    // to the event type string.
    unsigned.redacted_because = Some(EventJson::from(RedactionEvent {
        content: RedactionEventContent { reason: Some("redacted because".into()) },
        redacts: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        event_id: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        origin_server_ts: UNIX_EPOCH + Duration::from_millis(1),
        room_id: RoomId::try_from("!roomid:room.com").unwrap(),
        sender: UserId::try_from("@carl:example.com").unwrap(),
        unsigned: UnsignedData::default(),
    }));

    let redacted = json!({
      "event_id": "$h29iv0s8:example.com",
      "origin_server_ts": 1,
      "sender": "@carl:example.com",
      "unsigned": unsigned,
      "type": "m.room.message"
    });

    let actual = to_json_value(&redacted).unwrap();

    assert_matches!(
        from_json_value::<EventJson<AnySyncRoomEvent>>(actual)
            .unwrap()
            .deserialize()
            .unwrap(),
        AnySyncRoomEvent::RedactedMessage(AnyRedactedSyncMessageEvent::RoomMessage(RedactedSyncMessageEvent {
            event_id, content, ..
        })) if event_id == EventId::try_from("$h29iv0s8:example.com").unwrap()
            && is_zst(&content)
    )
}

#[test]
fn redacted_state_event_deserialize() {
    let mut unsigned = UnsignedData::default();
    unsigned.redacted_because = Some(EventJson::from(RedactionEvent {
        content: RedactionEventContent { reason: Some("redacted because".into()) },
        redacts: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        event_id: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        origin_server_ts: UNIX_EPOCH + Duration::from_millis(1),
        room_id: RoomId::try_from("!roomid:room.com").unwrap(),
        sender: UserId::try_from("@carl:example.com").unwrap(),
        unsigned: UnsignedData::default(),
    }));

    let redacted = json!({
      "content": {
        "creator": "@carl:example.com",
      },
      "event_id": "$h29iv0s8:example.com",
      "origin_server_ts": 1,
      "sender": "@carl:example.com",
      "state_key": "hello there",
      "unsigned": unsigned,
      "type": "m.room.create"
    });

    assert_matches!(
        from_json_value::<EventJson<AnySyncRoomEvent>>(redacted)
            .unwrap()
            .deserialize()
            .unwrap(),
        AnySyncRoomEvent::RedactedState(AnyRedactedSyncStateEvent::RoomCreate(RedactedSyncStateEvent {
            content: RedactedCreateEventContent {
                creator,
            },
            event_id, state_key, unsigned, ..
        })) if event_id == EventId::try_from("$h29iv0s8:example.com").unwrap()
            && unsigned.redacted_because.is_some()
            && state_key == "hello there"
            && creator == UserId::try_from("@carl:example.com").unwrap()
    )
}

#[test]
fn redacted_custom_event_serialize() {
    let mut unsigned = UnsignedData::default();
    unsigned.redacted_because = Some(EventJson::from(RedactionEvent {
        content: RedactionEventContent { reason: Some("redacted because".into()) },
        redacts: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        event_id: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        origin_server_ts: UNIX_EPOCH + Duration::from_millis(1),
        room_id: RoomId::try_from("!roomid:room.com").unwrap(),
        sender: UserId::try_from("@carl:example.com").unwrap(),
        unsigned: UnsignedData::default(),
    }));

    let redacted = json!({
        "event_id": "$h29iv0s8:example.com",
        "origin_server_ts": 1,
        "sender": "@carl:example.com",
        "state_key": "hello there",
        "unsigned": unsigned,
        "type": "m.made.up"
    });

    assert_matches!(
        from_json_value::<EventJson<AnySyncRoomEvent>>(redacted.clone())
            .unwrap()
            .deserialize()
            .unwrap(),
        AnySyncRoomEvent::RedactedState(AnyRedactedSyncStateEvent::Custom(RedactedSyncStateEvent {
            content: RedactedCustomEventContent {
                event_type,
            },
            event_id, state_key, unsigned, ..
        })) if event_id == EventId::try_from("$h29iv0s8:example.com").unwrap()
            && unsigned.redacted_because.is_some()
            && state_key == "hello there"
            && event_type == "m.made.up"
    );

    let x = from_json_value::<EventJson<crate::AnyRedactedSyncStateEvent>>(redacted)
        .unwrap()
        .deserialize()
        .unwrap();
    assert_eq!(x.event_id(), &EventId::try_from("$h29iv0s8:example.com").unwrap())
}

#[test]
fn redacted_custom_event_deserialize() {
    let mut unsigned = UnsignedData::default();
    unsigned.redacted_because = Some(EventJson::from(RedactionEvent {
        content: RedactionEventContent { reason: Some("redacted because".into()) },
        redacts: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        event_id: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        origin_server_ts: UNIX_EPOCH + Duration::from_millis(1),
        room_id: RoomId::try_from("!roomid:room.com").unwrap(),
        sender: UserId::try_from("@carl:example.com").unwrap(),
        unsigned: UnsignedData::default(),
    }));

    let redacted = RedactedSyncStateEvent {
        content: RedactedCustomEventContent { event_type: "m.made.up".into() },
        event_id: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        sender: UserId::try_from("@carl:example.com").unwrap(),
        state_key: "hello there".into(),
        origin_server_ts: UNIX_EPOCH + Duration::from_millis(1),
        unsigned: unsigned.clone(),
    };

    let expected = json!({
      "event_id": "$h29iv0s8:example.com",
      "origin_server_ts": 1,
      "sender": "@carl:example.com",
      "state_key": "hello there",
      "unsigned": unsigned,
      "type": "m.made.up"
    });

    let actual = to_json_value(&redacted).unwrap();
    assert_eq!(actual, expected);
}
<commit_msg>Move UnsignedData creation into function for redacted tests<commit_after>use std::{
    convert::TryFrom,
    time::{Duration, UNIX_EPOCH},
};

use matches::assert_matches;
use ruma_events::{
    custom::RedactedCustomEventContent,
    room::{
        aliases::RedactedAliasesEventContent,
        create::RedactedCreateEventContent,
        message::RedactedMessageEventContent,
        redaction::{RedactionEvent, RedactionEventContent},
    },
    AnyRedactedMessageEvent, AnyRedactedSyncMessageEvent, AnyRedactedSyncStateEvent, AnyRoomEvent,
    AnySyncRoomEvent, EventJson, RedactedMessageEvent, RedactedSyncMessageEvent,
    RedactedSyncStateEvent, UnsignedData,
};
use ruma_identifiers::{EventId, RoomId, UserId};
use serde_json::{from_value as from_json_value, json, to_value as to_json_value};

fn is_zst<T>(_: &T) -> bool {
    std::mem::size_of::<T>() == 0
}

fn full_unsigned() -> UnsignedData {
    let mut unsigned = UnsignedData::default();
    // The presence of `redacted_because` triggers the event enum to return early
    // with `RedactedContent` instead of failing to deserialize according
    // to the event type string.
    unsigned.redacted_because = Some(EventJson::from(RedactionEvent {
        content: RedactionEventContent { reason: Some("redacted because".into()) },
        redacts: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        event_id: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        origin_server_ts: UNIX_EPOCH + Duration::from_millis(1),
        room_id: RoomId::try_from("!roomid:room.com").unwrap(),
        sender: UserId::try_from("@carl:example.com").unwrap(),
        unsigned: UnsignedData::default(),
    }));

    unsigned
}

#[test]
fn redacted_message_event_serialize() {
    let redacted = RedactedSyncMessageEvent {
        content: RedactedMessageEventContent,
        event_id: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        origin_server_ts: UNIX_EPOCH + Duration::from_millis(1),
        sender: UserId::try_from("@carl:example.com").unwrap(),
        unsigned: UnsignedData::default(),
    };

    let expected = json!({
      "event_id": "$h29iv0s8:example.com",
      "origin_server_ts": 1,
      "sender": "@carl:example.com",
      "type": "m.room.message"
    });

    let actual = to_json_value(&redacted).unwrap();
    assert_eq!(actual, expected);
}

#[test]
fn redacted_aliases_event_serialize() {
    let redacted = RedactedSyncStateEvent {
        content: RedactedAliasesEventContent { aliases: None },
        event_id: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        state_key: "".into(),
        origin_server_ts: UNIX_EPOCH + Duration::from_millis(1),
        sender: UserId::try_from("@carl:example.com").unwrap(),
        unsigned: UnsignedData::default(),
    };

    let expected = json!({
      "event_id": "$h29iv0s8:example.com",
      "state_key": "",
      "origin_server_ts": 1,
      "sender": "@carl:example.com",
      "type": "m.room.aliases"
    });

    let actual = to_json_value(&redacted).unwrap();
    assert_eq!(actual, expected);
}

#[test]
fn redacted_deserialize_any_room() {
    let unsigned = full_unsigned();

    let redacted = json!({
      "event_id": "$h29iv0s8:example.com",
      "room_id": "!roomid:room.com",
      "origin_server_ts": 1,
      "sender": "@carl:example.com",
      "unsigned": unsigned,
      "type": "m.room.message"
    });

    let actual = to_json_value(&redacted).unwrap();

    assert_matches!(
        from_json_value::<EventJson<AnyRoomEvent>>(actual)
            .unwrap()
            .deserialize()
            .unwrap(),
        AnyRoomEvent::RedactedMessage(AnyRedactedMessageEvent::RoomMessage(RedactedMessageEvent {
            event_id, room_id, content, ..
        })) if event_id == EventId::try_from("$h29iv0s8:example.com").unwrap()
            && room_id == RoomId::try_from("!roomid:room.com").unwrap()
            && is_zst(&content)
    )
}

#[test]
fn redacted_deserialize_any_room_sync() {
    let mut unsigned = UnsignedData::default();
    // The presence of `redacted_because` triggers the event enum (AnySyncRoomEvent in this case)
    // to return early with `RedactedContent` instead of failing to deserialize according
    // to the event type string.
    unsigned.redacted_because = Some(EventJson::from(RedactionEvent {
        content: RedactionEventContent { reason: Some("redacted because".into()) },
        redacts: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        event_id: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        origin_server_ts: UNIX_EPOCH + Duration::from_millis(1),
        room_id: RoomId::try_from("!roomid:room.com").unwrap(),
        sender: UserId::try_from("@carl:example.com").unwrap(),
        unsigned: UnsignedData::default(),
    }));

    let redacted = json!({
      "event_id": "$h29iv0s8:example.com",
      "origin_server_ts": 1,
      "sender": "@carl:example.com",
      "unsigned": unsigned,
      "type": "m.room.message"
    });

    let actual = to_json_value(&redacted).unwrap();

    assert_matches!(
        from_json_value::<EventJson<AnySyncRoomEvent>>(actual)
            .unwrap()
            .deserialize()
            .unwrap(),
        AnySyncRoomEvent::RedactedMessage(AnyRedactedSyncMessageEvent::RoomMessage(RedactedSyncMessageEvent {
            event_id, content, ..
        })) if event_id == EventId::try_from("$h29iv0s8:example.com").unwrap()
            && is_zst(&content)
    )
}

#[test]
fn redacted_state_event_deserialize() {
    let unsigned = full_unsigned();

    let redacted = json!({
      "content": {
        "creator": "@carl:example.com",
      },
      "event_id": "$h29iv0s8:example.com",
      "origin_server_ts": 1,
      "sender": "@carl:example.com",
      "state_key": "hello there",
      "unsigned": unsigned,
      "type": "m.room.create"
    });

    assert_matches!(
        from_json_value::<EventJson<AnySyncRoomEvent>>(redacted)
            .unwrap()
            .deserialize()
            .unwrap(),
        AnySyncRoomEvent::RedactedState(AnyRedactedSyncStateEvent::RoomCreate(RedactedSyncStateEvent {
            content: RedactedCreateEventContent {
                creator,
            },
            event_id, state_key, unsigned, ..
        })) if event_id == EventId::try_from("$h29iv0s8:example.com").unwrap()
            && unsigned.redacted_because.is_some()
            && state_key == "hello there"
            && creator == UserId::try_from("@carl:example.com").unwrap()
    )
}

#[test]
fn redacted_custom_event_serialize() {
    let unsigned = full_unsigned();

    let redacted = json!({
        "event_id": "$h29iv0s8:example.com",
        "origin_server_ts": 1,
        "sender": "@carl:example.com",
        "state_key": "hello there",
        "unsigned": unsigned,
        "type": "m.made.up"
    });

    assert_matches!(
        from_json_value::<EventJson<AnySyncRoomEvent>>(redacted.clone())
            .unwrap()
            .deserialize()
            .unwrap(),
        AnySyncRoomEvent::RedactedState(AnyRedactedSyncStateEvent::Custom(RedactedSyncStateEvent {
            content: RedactedCustomEventContent {
                event_type,
            },
            event_id, state_key, unsigned, ..
        })) if event_id == EventId::try_from("$h29iv0s8:example.com").unwrap()
            && unsigned.redacted_because.is_some()
            && state_key == "hello there"
            && event_type == "m.made.up"
    );

    let x = from_json_value::<EventJson<crate::AnyRedactedSyncStateEvent>>(redacted)
        .unwrap()
        .deserialize()
        .unwrap();
    assert_eq!(x.event_id(), &EventId::try_from("$h29iv0s8:example.com").unwrap())
}

#[test]
fn redacted_custom_event_deserialize() {
    let unsigned = full_unsigned();

    let redacted = RedactedSyncStateEvent {
        content: RedactedCustomEventContent { event_type: "m.made.up".into() },
        event_id: EventId::try_from("$h29iv0s8:example.com").unwrap(),
        sender: UserId::try_from("@carl:example.com").unwrap(),
        state_key: "hello there".into(),
        origin_server_ts: UNIX_EPOCH + Duration::from_millis(1),
        unsigned: unsigned.clone(),
    };

    let expected = json!({
      "event_id": "$h29iv0s8:example.com",
      "origin_server_ts": 1,
      "sender": "@carl:example.com",
      "state_key": "hello there",
      "unsigned": unsigned,
      "type": "m.made.up"
    });

    let actual = to_json_value(&redacted).unwrap();
    assert_eq!(actual, expected);
}
<|endoftext|>


<commit_before>#![feature(core_intrinsics)]
use std::fs::File;
use std::io::{BufRead, BufReader};

use std::intrinsics::unlikely;

const VALUE: [u64; 256] = {
    let mut array = [0u64; 256];
    array[b'A' as usize] = 1;
    array[b'T' as usize] = 1;
    array[b'G' as usize] = 1 << 32;
    array[b'C' as usize] = 1 << 32;
    array
};

fn main() {
    let filename = "chry_multiplied.fa";

    let file = File::open(filename).unwrap();
    let mut reader = BufReader::new(file);

    let mut totals = 0u64;
    let mut line = Vec::with_capacity(100);
    loop {
        if unlikely(reader.read_until(b'\n', &mut line).expect("error reading") == 0) {
            break;
        }

        if unlikely(line[0] == b'>') {
            line.clear();
            continue;
        }

        for c in line.iter() {
            totals += VALUE[*c as usize];
        }
        line.clear();
    }
    let at = totals & 0xFFFFFFFF;
    let gc = totals >> 32;
    let gc_ratio: f64 = gc as f64 / (gc as f64 + at as f64);
    println!("GC ratio (gc/(gc+at)): {}", gc_ratio);
}
<commit_msg>Give the compiler a hit to unroll or vectorize<commit_after>#![feature(core_intrinsics)]
use std::fs::File;
use std::io::{BufRead, BufReader};

use std::intrinsics::unlikely;

const VALUE: [u64; 256] = {
    let mut array = [0u64; 256];
    array[b'A' as usize] = 1;
    array[b'T' as usize] = 1;
    array[b'G' as usize] = 1 << 32;
    array[b'C' as usize] = 1 << 32;
    array
};

const CHUNKSIZE: usize = 16;

fn main() {
    let filename = "chry_multiplied.fa";

    let file = File::open(filename).unwrap();
    let mut reader = BufReader::new(file);

    let mut totals = 0u64;
    let mut line = Vec::with_capacity(100);
    loop {
        if unlikely(reader.read_until(b'\n', &mut line).expect("error reading") == 0) {
            break;
        }

        if unlikely(line[0] == b'>') {
            line.clear();
            continue;
        }
        let mut chunker = line.chunks_exact(CHUNKSIZE);
        while let Some(chars) = chunker.next() {
            totals += (0..CHUNKSIZE)
                .into_iter()
                .map(|c| chars[c])
                .map(|b| VALUE[b as usize])
                .sum::<u64>();
        }
        for c in chunker.remainder() {
            totals += VALUE[*c as usize];
        }
        line.clear();
    }
    let at = totals & 0xFFFFFFFF;
    let gc = totals >> 32;
    let gc_ratio: f64 = gc as f64 / (gc as f64 + at as f64);
    println!("GC ratio (gc/(gc+at)): {}", gc_ratio);
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::io;
use std::ops::Neg;
use std::os::unix::process::{CommandExt, ExitStatusExt};
use std::process::{Command, ExitStatus, Stdio};
use std::result;

use core::os;
use core::os::process::{Pid, signal, Signal};
use libc::{self, c_int, pid_t};
use protocol::{self, ShutdownMethod};
use time::{Duration, SteadyTime};

use error::{Error, Result};
use service::Service;

pub struct Process {
    pid: pid_t,
    status: Option<ExitStatus>,
}

impl Process {
    fn new(pid: u32) -> Self {
        Process {
            pid: pid as pid_t,
            status: None,
        }
    }

    pub fn id(&self) -> Pid {
        self.pid
    }

    /// Attempt to gracefully terminate a proccess and then forcefully kill it after
    /// 8 seconds if it has not terminated.
    pub fn kill(&mut self) -> ShutdownMethod {
        // check the group of the process being killed
        // if it is the root process of the process group
        // we send our signals to the entire process group
        // to prevent orphaned processes.
        let pgid = unsafe { libc::getpgid(self.pid) };
        if self.pid == pgid {
            debug!(
                "pid to kill {} is the process group root. Sending signal to process group.",
                self.pid
            );
            // sending a signal to the negative pid sends it to the
            // entire process group instead just the single pid
            self.pid = self.pid.neg();
        }

        // JW TODO: Determine if the error represents a case where the process was already
        // exited before we return out and assume so.
        if signal(self.id(), Signal::TERM).is_err() {
            return ShutdownMethod::AlreadyExited;
        }
        let stop_time = SteadyTime::now() + Duration::seconds(8);
        loop {
            if let Ok(Some(_status)) = self.try_wait() {
                return ShutdownMethod::GracefulTermination;
            }
            if SteadyTime::now() < stop_time {
                continue;
            }
            // JW TODO: Determine if the error represents a case where the process was already
            // exited before we return out and assume so.
            if signal(self.id(), Signal::KILL).is_err() {
                return ShutdownMethod::GracefulTermination;
            }
            return ShutdownMethod::Killed;
        }
    }

    pub fn try_wait(&mut self) -> Result<Option<ExitStatus>> {
        if let Some(status) = self.status {
            return Ok(Some(status));
        }
        let mut status = 0 as c_int;
        match unsafe { libc::waitpid(self.pid, &mut status, libc::WNOHANG) } {
            0 => Ok(None),
            -1 => Err(Error::ExecWait(io::Error::last_os_error())),
            _ => {
                self.status = Some(ExitStatus::from_raw(status));
                Ok(Some(ExitStatus::from_raw(status)))
            }
        }
    }

    pub fn wait(&mut self) -> Result<ExitStatus> {
        if let Some(status) = self.status {
            return Ok(status);
        }
        let mut status = 0 as c_int;
        match unsafe { libc::waitpid(self.pid, &mut status, 0) } {
            -1 => Err(Error::ExecWait(io::Error::last_os_error())),
            _ => {
                self.status = Some(ExitStatus::from_raw(status));
                Ok(ExitStatus::from_raw(status))
            }
        }
    }
}

pub fn run(msg: protocol::Spawn) -> Result<Service> {
    debug!("launcher is spawning {}", msg.get_binary());
    let mut cmd = Command::new(msg.get_binary());
    let uid = os::users::get_uid_by_name(msg.get_svc_user()).ok_or(
        Error::UserNotFound(msg.get_svc_user().to_string()),
    )?;
    let gid = os::users::get_gid_by_name(msg.get_svc_group()).ok_or(
        Error::GroupNotFound(msg.get_svc_group().to_string()),
    )?;
    cmd.before_exec(owned_pgid);
    cmd.stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .uid(uid)
        .gid(gid);
    for (key, val) in msg.get_env().iter() {
        cmd.env(key, val);
    }
    let child = cmd.spawn().map_err(Error::Spawn)?;
    let process = Process::new(child.id());
    Ok(Service::new(msg, process, child.stdout, child.stderr))
}

// we want the command to spawn processes in their own process group
// and not the same group as the Launcher. Otherwise if a child process
// sends SIGTERM to the group, the Launcher could be terminated.
fn owned_pgid() -> result::Result<(), io::Error> {
    unsafe {
        libc::setpgid(0, 0);
    }
    Ok(())
}
<commit_msg>when killing a process group, dont mutate the pid of the process struct<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::io;
use std::ops::Neg;
use std::os::unix::process::{CommandExt, ExitStatusExt};
use std::process::{Command, ExitStatus, Stdio};
use std::result;

use core::os;
use core::os::process::{Pid, signal, Signal};
use libc::{self, c_int, pid_t};
use protocol::{self, ShutdownMethod};
use time::{Duration, SteadyTime};

use error::{Error, Result};
use service::Service;

pub struct Process {
    pid: pid_t,
    status: Option<ExitStatus>,
}

impl Process {
    fn new(pid: u32) -> Self {
        Process {
            pid: pid as pid_t,
            status: None,
        }
    }

    pub fn id(&self) -> Pid {
        self.pid
    }

    /// Attempt to gracefully terminate a proccess and then forcefully kill it after
    /// 8 seconds if it has not terminated.
    pub fn kill(&mut self) -> ShutdownMethod {
        let mut pid_to_kill = self.pid;
        // check the group of the process being killed
        // if it is the root process of the process group
        // we send our signals to the entire process group
        // to prevent orphaned processes.
        let pgid = unsafe { libc::getpgid(self.pid) };
        if self.pid == pgid {
            debug!(
                "pid to kill {} is the process group root. Sending signal to process group.",
                self.pid
            );
            // sending a signal to the negative pid sends it to the
            // entire process group instead just the single pid
            pid_to_kill = self.pid.neg();
        }

        // JW TODO: Determine if the error represents a case where the process was already
        // exited before we return out and assume so.
        if signal(pid_to_kill, Signal::TERM).is_err() {
            return ShutdownMethod::AlreadyExited;
        }
        let stop_time = SteadyTime::now() + Duration::seconds(8);
        loop {
            if let Ok(Some(_status)) = self.try_wait() {
                return ShutdownMethod::GracefulTermination;
            }
            if SteadyTime::now() < stop_time {
                continue;
            }
            // JW TODO: Determine if the error represents a case where the process was already
            // exited before we return out and assume so.
            if signal(pid_to_kill, Signal::KILL).is_err() {
                return ShutdownMethod::GracefulTermination;
            }
            return ShutdownMethod::Killed;
        }
    }

    pub fn try_wait(&mut self) -> Result<Option<ExitStatus>> {
        if let Some(status) = self.status {
            return Ok(Some(status));
        }
        let mut status = 0 as c_int;
        match unsafe { libc::waitpid(self.pid, &mut status, libc::WNOHANG) } {
            0 => Ok(None),
            -1 => Err(Error::ExecWait(io::Error::last_os_error())),
            _ => {
                self.status = Some(ExitStatus::from_raw(status));
                Ok(Some(ExitStatus::from_raw(status)))
            }
        }
    }

    pub fn wait(&mut self) -> Result<ExitStatus> {
        if let Some(status) = self.status {
            return Ok(status);
        }
        let mut status = 0 as c_int;
        match unsafe { libc::waitpid(self.pid, &mut status, 0) } {
            -1 => Err(Error::ExecWait(io::Error::last_os_error())),
            _ => {
                self.status = Some(ExitStatus::from_raw(status));
                Ok(ExitStatus::from_raw(status))
            }
        }
    }
}

pub fn run(msg: protocol::Spawn) -> Result<Service> {
    debug!("launcher is spawning {}", msg.get_binary());
    let mut cmd = Command::new(msg.get_binary());
    let uid = os::users::get_uid_by_name(msg.get_svc_user()).ok_or(
        Error::UserNotFound(msg.get_svc_user().to_string()),
    )?;
    let gid = os::users::get_gid_by_name(msg.get_svc_group()).ok_or(
        Error::GroupNotFound(msg.get_svc_group().to_string()),
    )?;
    cmd.before_exec(owned_pgid);
    cmd.stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .uid(uid)
        .gid(gid);
    for (key, val) in msg.get_env().iter() {
        cmd.env(key, val);
    }
    let child = cmd.spawn().map_err(Error::Spawn)?;
    let process = Process::new(child.id());
    Ok(Service::new(msg, process, child.stdout, child.stderr))
}

// we want the command to spawn processes in their own process group
// and not the same group as the Launcher. Otherwise if a child process
// sends SIGTERM to the group, the Launcher could be terminated.
fn owned_pgid() -> result::Result<(), io::Error> {
    unsafe {
        libc::setpgid(0, 0);
    }
    Ok(())
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::sync::Arc;

use common::{PerfLogger, SourceLocationKey};
use graphql_ir::{
    build_ir_with_extra_features, BuilderOptions, FragmentVariablesSemantic, OperationDefinition,
    Program, Selection,
};
use graphql_syntax::{parse_executable_with_error_recovery, ExecutableDefinition};
use graphql_text_printer::print_full_operation;
use interner::{Intern, StringKey};
use lsp_types::request::Request;
use relay_compiler::{
    apply_transforms,
    config::{Config, ProjectConfig},
};
use serde_json::Value;

use crate::{
    lsp_extra_data_provider::Query,
    lsp_runtime_error::LSPRuntimeResult,
    server::{LSPState, SourcePrograms},
    Schemas,
};
use serde::{Deserialize, Serialize};

pub(crate) enum GraphQlExecuteQuery {}

#[derive(Serialize, Deserialize)]
pub(crate) struct GraphQlExecuteQueryParams {
    text: String,
    variables: serde_json::Value,
}

#[derive(Serialize, Deserialize)]
pub(crate) enum GraphQlResponse {
    #[serde(rename = "data")]
    Data(serde_json::Value),
    #[serde(rename = "error")]
    Error(serde_json::Value),
}

impl Request for GraphQlExecuteQuery {
    type Params = GraphQlExecuteQueryParams;
    type Result = GraphQlResponse;
    const METHOD: &'static str = "graphql/executeQuery";
}

/// This function will return the program that contains only operation
/// and all referenced fragments.
/// We can use it to print the full query text
fn get_operation_only_program(
    operation: Arc<OperationDefinition>,
    project_name: StringKey,
    programs: &SourcePrograms,
) -> Option<Program> {
    let program = programs.get(&project_name)?;

    let mut selections_to_visit: Vec<_> = vec![&operation.selections];
    let mut next_program = Program::new(program.schema.clone());
    next_program.insert_operation(Arc::clone(&operation));

    while !selections_to_visit.is_empty() {
        let current_selections = selections_to_visit.pop()?;
        for selection in current_selections {
            match selection {
                Selection::FragmentSpread(spread) => {
                    let fragment = program
                        .fragment(spread.fragment.item)
                        .expect("Expect fragment to exist");
                    selections_to_visit.push(&fragment.selections);
                    next_program.insert_fragment(Arc::clone(fragment));
                }
                Selection::Condition(condition) => {
                    selections_to_visit.push(&condition.selections);
                }
                Selection::LinkedField(linked_field) => {
                    selections_to_visit.push(&linked_field.selections);
                }
                Selection::InlineFragment(inline_fragment) => {
                    selections_to_visit.push(&inline_fragment.selections);
                }
                Selection::ScalarField(_) => {}
            }
        }
    }

    Some(next_program)
}

/// Given the `Program` that contain operation+all its fragments this
/// function will `apply_transforms` and return a query text,
/// that can be sent to the server
fn get_full_operation_text<TPerfLogger: PerfLogger + 'static>(
    project_name: StringKey,
    project_config: &ProjectConfig,
    config: Arc<Config>,
    operation_name: StringKey,
    source_program: Arc<Program>,
    perf_logger: Arc<TPerfLogger>,
) -> String {
    let programs = apply_transforms(
        project_name,
        source_program,
        Default::default(),
        &config.connection_interface,
        Arc::new(
            project_config
                .feature_flags
                .as_ref()
                .cloned()
                .unwrap_or_else(|| config.feature_flags.clone()),
        ),
        perf_logger,
    )
    .unwrap();

    let print_operation_node = programs
        .operation_text
        .operation(operation_name)
        .expect("a query text operation should be generated for this operation");

    print_full_operation(&programs.operation_text, print_operation_node)
}

/// Given only one operation AST it will construct the typed IR
/// But it will ignore missing fragment spreads
fn build_operation_ir(
    project_config: &ProjectConfig,
    schemas: &Schemas,
    operation: graphql_syntax::OperationDefinition,
    operation_name: StringKey,
) -> Result<OperationDefinition, String> {
    let schema = schemas.get(&project_config.name).unwrap().clone();
    let definitions = vec![graphql_syntax::ExecutableDefinition::Operation(operation)];
    let ir = build_ir_with_extra_features(
        &schema,
        &definitions,
        BuilderOptions {
            allow_undefined_fragment_spreads: true,
            fragment_variables_semantic: FragmentVariablesSemantic::PassedValue,
            relay_mode: true,
        },
    )
    .map_err(|errors| format!("{:?}", errors))?;

    for item in ir {
        if let graphql_ir::ExecutableDefinition::Operation(op) = item {
            if op.name.item == operation_name {
                return Ok(op);
            }
        }
    }

    Err("Operation could not be constructed".to_string())
}

pub(crate) fn on_graphql_execute_query<TPerfLogger: PerfLogger + 'static>(
    state: &mut LSPState<TPerfLogger>,
    params: GraphQlExecuteQueryParams,
) -> LSPRuntimeResult<<GraphQlExecuteQuery as Request>::Result> {
    let result = parse_executable_with_error_recovery(&params.text, SourceLocationKey::Generated);
    // TODO: Figure out the project form request params
    let project_name = "facebook".intern();

    if !&result.errors.is_empty() {
        return Ok(GraphQlResponse::Error(Value::Array(
            result
                .errors
                .iter()
                .map(|err| Value::String(format!("{:?}", err)))
                .collect::<Vec<_>>(),
        )));
    }

    let query_text: String = match &result.item.definitions[0] {
        ExecutableDefinition::Fragment(_) => {
            return Ok(GraphQlResponse::Error(Value::String(
                "Cannot fetch data for fragments".to_string(),
            )));
        }
        ExecutableDefinition::Operation(operation) => match &operation.name {
            Some(operation_name) => {
                let project_config = state.get_project_config_ref(project_name).unwrap();

                let operation_ir = build_operation_ir(
                    project_config,
                    &state.get_schemas(),
                    operation.clone(),
                    operation_name.value,
                )
                .unwrap();

                let operation_name = operation_ir.name.item;

                // Internally, we have a rule that only operations with names can be used
                if let Some(operation_sources) = get_operation_only_program(
                    Arc::new(operation_ir),
                    project_name,
                    state.get_source_programs_ref(),
                ) {
                    let operation_text = get_full_operation_text(
                        project_name,
                        state.get_project_config_ref(project_name).unwrap(),
                        state.get_config(),
                        operation_name,
                        Arc::new(operation_sources),
                        state.get_logger(),
                    );

                    operation_text
                } else {
                    // If, for some reason we could not get the operation text from the sources
                    // we will send the original text to the server
                    params.text
                }
            }
            None => params.text,
        },
    };

    match state
        .extra_data_provider
        .execute_query(Query::Text(query_text), params.variables)
    {
        Ok(data) => Ok(GraphQlResponse::Data(data)),
        Err(error) => Ok(GraphQlResponse::Error(error)),
    }
}
<commit_msg>Create a program for printing with queries, and fragment from the initial text<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::{collections::HashSet, sync::Arc};

use common::{PerfLogger, SourceLocationKey};
use graphql_ir::{
    build_ir_with_extra_features, BuilderOptions, ExecutableDefinition, FragmentDefinition,
    FragmentVariablesSemantic, OperationDefinition, Program, Selection,
};
use graphql_syntax::parse_executable_with_error_recovery;
use graphql_text_printer::print_full_operation;
use interner::{Intern, StringKey};
use lsp_types::request::Request;
use relay_compiler::{
    apply_transforms,
    config::{Config, ProjectConfig},
    Programs,
};
use schema::SDLSchema;
use serde_json::Value;

use crate::{
    lsp_extra_data_provider::Query,
    lsp_runtime_error::LSPRuntimeResult,
    server::{LSPState, SourcePrograms},
};
use serde::{Deserialize, Serialize};

pub(crate) enum GraphQlExecuteQuery {}

#[derive(Serialize, Deserialize)]
pub(crate) struct GraphQlExecuteQueryParams {
    text: String,
    variables: serde_json::Value,
}

#[derive(Serialize, Deserialize)]
pub(crate) enum GraphQlResponse {
    #[serde(rename = "data")]
    Data(serde_json::Value),
    #[serde(rename = "error")]
    Error(serde_json::Value),
}

impl Request for GraphQlExecuteQuery {
    type Params = GraphQlExecuteQueryParams;
    type Result = GraphQlResponse;
    const METHOD: &'static str = "graphql/executeQuery";
}

/// This function will return the program that contains only operation
/// and all referenced fragments.
/// We can use it to print the full query text
fn get_operation_only_program(
    operation: Arc<OperationDefinition>,
    fragments: Vec<Arc<FragmentDefinition>>,
    project_name: StringKey,
    programs: &SourcePrograms,
) -> Option<Program> {
    let full_program = programs.get(&project_name)?;

    let mut selections_to_visit: Vec<_> = vec![&operation.selections];
    let mut next_program = Program::new(full_program.schema.clone());
    next_program.insert_operation(Arc::clone(&operation));
    for fragment in fragments.iter() {
        selections_to_visit.push(&fragment.selections);
        next_program.insert_fragment(Arc::clone(fragment));
    }

    let mut visited_fragments: HashSet<StringKey> = HashSet::default();

    while !selections_to_visit.is_empty() {
        let current_selections = selections_to_visit.pop()?;
        for selection in current_selections {
            match selection {
                Selection::FragmentSpread(spread) => {
                    // Skip, if we already visited this fragment
                    if visited_fragments.contains(&spread.fragment.item) {
                        continue;
                    }
                    visited_fragments.insert(spread.fragment.item);
                    // Second, if this fragment is already in the `next_program`,
                    // it selection already added to the visiting stack
                    if next_program.fragment(spread.fragment.item).is_some() {
                        continue;
                    }

                    // Finally, add all missing fragment spreads from the full program
                    let fragment = full_program
                        .fragment(spread.fragment.item)
                        .expect("Expect fragment to exist");
                    selections_to_visit.push(&fragment.selections);
                    next_program.insert_fragment(Arc::clone(fragment));
                }
                Selection::Condition(condition) => {
                    selections_to_visit.push(&condition.selections);
                }
                Selection::LinkedField(linked_field) => {
                    selections_to_visit.push(&linked_field.selections);
                }
                Selection::InlineFragment(inline_fragment) => {
                    selections_to_visit.push(&inline_fragment.selections);
                }
                Selection::ScalarField(_) => {}
            }
        }
    }

    Some(next_program)
}

/// Given the `Program` that contain operation+all its fragments this
/// function will `apply_transforms` and create a full set of programs, including the one
/// that may generate full operation text
fn transform_program<TPerfLogger: PerfLogger + 'static>(
    project_config: &ProjectConfig,
    config: Arc<Config>,
    program: Arc<Program>,
    perf_logger: Arc<TPerfLogger>,
) -> Result<Programs, String> {
    apply_transforms(
        project_config.name,
        program,
        Default::default(),
        &config.connection_interface,
        Arc::new(
            project_config
                .feature_flags
                .as_ref()
                .cloned()
                .unwrap_or_else(|| config.feature_flags.clone()),
        ),
        perf_logger,
    )
    .map_err(|err| format!("{:?}", err))
}

fn print_full_operation_text(programs: Programs, operation_name: StringKey) -> String {
    let print_operation_node = programs
        .operation_text
        .operation(operation_name)
        .expect("a query text operation should be generated for this operation");

    print_full_operation(&programs.operation_text, print_operation_node)
}

/// From the list of AST nodes we're trying to extract the operation and possible
/// list of fragment, to construct the initial `Program` that we could later transform
/// and print
fn build_operation_ir_with_fragments(
    definitions: &[graphql_syntax::ExecutableDefinition],
    schema: Arc<SDLSchema>,
) -> Result<(Arc<OperationDefinition>, Vec<Arc<FragmentDefinition>>), String> {
    let ir = build_ir_with_extra_features(
        &schema,
        definitions,
        BuilderOptions {
            allow_undefined_fragment_spreads: true,
            fragment_variables_semantic: FragmentVariablesSemantic::PassedValue,
            relay_mode: true,
        },
    )
    .map_err(|errors| format!("{:?}", errors))?;

    if let Some(operation) = ir.iter().find_map(|item| {
        if let ExecutableDefinition::Operation(operation) = item {
            Some(Arc::new(operation.clone()))
        } else {
            None
        }
    }) {
        let fragments = ir
            .iter()
            .filter_map(|item| {
                if let ExecutableDefinition::Fragment(fragment) = item {
                    Some(Arc::new(fragment.clone()))
                } else {
                    None
                }
            })
            .collect::<Vec<_>>();

        Ok((operation, fragments))
    } else {
        Err("Unable to find an operation.".to_string())
    }
}

pub(crate) fn on_graphql_execute_query<TPerfLogger: PerfLogger + 'static>(
    state: &mut LSPState<TPerfLogger>,
    params: GraphQlExecuteQueryParams,
) -> LSPRuntimeResult<<GraphQlExecuteQuery as Request>::Result> {
    let result = parse_executable_with_error_recovery(&params.text, SourceLocationKey::Generated);
    // TODO: Figure out the project form request params
    let project_name = "facebook".intern();

    if !&result.errors.is_empty() {
        return Ok(GraphQlResponse::Error(Value::Array(
            result
                .errors
                .iter()
                .map(|err| Value::String(format!("{:?}", err)))
                .collect::<Vec<_>>(),
        )));
    }

    let project_config = state.get_project_config_ref(project_name).unwrap();
    let schema = state
        .get_schemas()
        .get(&project_config.name)
        .unwrap()
        .clone();

    let (operation, fragments) =
        build_operation_ir_with_fragments(&result.item.definitions, schema).unwrap();

    let operation_name = operation.name.item;
    let query_text = if let Some(operation_program) = get_operation_only_program(
        operation,
        fragments,
        project_name,
        state.get_source_programs_ref(),
    ) {
        let programs = transform_program(
            project_config,
            state.get_config(),
            Arc::new(operation_program),
            state.get_logger(),
        )
        .unwrap();

        print_full_operation_text(programs, operation_name)
    } else {
        // If, for some reason we could not get the operation text from the sources
        // we will send the original text to the server
        params.text
    };

    match state
        .extra_data_provider
        .execute_query(Query::Text(query_text), params.variables)
    {
        Ok(data) => Ok(GraphQlResponse::Data(data)),
        Err(error) => Ok(GraphQlResponse::Error(error)),
    }
}
<|endoftext|>