<commit_before>// Copyright 2021 Google LLC
//
// 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
//
//     https://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.
//
// Also licensed under MIT license, at your choice.

//! A simple compute shader example that draws into a window, based on wgpu.

use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Instant;

use rust_gpu_toy_shared::Config;
use wgpu::util::make_spirv;
use wgpu::{Extent3d, ShaderModule};

use winit::dpi::PhysicalSize;
use winit::{
    event::{Event, WindowEvent},
    event_loop::{ControlFlow, EventLoop},
    window::Window,
};

use spirv_builder::{CompileResult, SpirvBuilder};

struct Pipelines {
    compute_pipeline: wgpu::ComputePipeline,

    copy_pipeline: wgpu::RenderPipeline,
}

struct Layouts {
    compute_bg_layout: wgpu::BindGroupLayout,
    copy_bg_layout: wgpu::BindGroupLayout,

    compute_pipeline_layout: wgpu::PipelineLayout,
    copy_pipeline_layout: wgpu::PipelineLayout,
}

/// State of the runner, inspired by
/// https://sotrh.github.io/learn-wgpu/
struct State {
    pipelines: Pipelines,
    layouts: Layouts,

    surface: wgpu::Surface,
    device: wgpu::Device,
    queue: wgpu::Queue,

    sc_desc: wgpu::SwapChainDescriptor,
    swap_chain: wgpu::SwapChain,
    size: winit::dpi::PhysicalSize<u32>,

    sampler: wgpu::Sampler,

    compute_bind_group: wgpu::BindGroup,
    copy_bind_group: wgpu::BindGroup,

    start_time: Instant,
}

impl State {
    async fn new(window: &Window, compilation: CompileResult) -> Self {
        let size = window.inner_size();
        let instance = wgpu::Instance::new(wgpu::BackendBit::PRIMARY);
        let surface = unsafe { instance.create_surface(window) };
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: Default::default(),
                compatible_surface: Some(&surface),
            })
            .await
            .expect("error finding adapter");
        let features = wgpu::Features::PUSH_CONSTANTS;
        let limits = wgpu::Limits {
            max_push_constant_size: std::mem::size_of::<Config>() as u32,
            ..Default::default()
        };
        let (device, queue) = adapter
            .request_device(
                &wgpu::DeviceDescriptor {
                    label: None,
                    features,
                    limits,
                },
                None,
            )
            .await
            .expect("error creating device");
        let swapchain_format = adapter.get_swap_chain_preferred_format(&surface).unwrap();
        let sc_desc = wgpu::SwapChainDescriptor {
            usage: wgpu::TextureUsage::RENDER_ATTACHMENT,
            format: swapchain_format,
            width: size.width,
            height: size.height,
            present_mode: wgpu::PresentMode::Mailbox,
        };

        let swap_chain = device.create_swap_chain(&surface, &sc_desc);

        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            address_mode_u: wgpu::AddressMode::ClampToEdge,
            address_mode_v: wgpu::AddressMode::ClampToEdge,
            address_mode_w: wgpu::AddressMode::ClampToEdge,
            mag_filter: wgpu::FilterMode::Nearest,
            min_filter: wgpu::FilterMode::Nearest,
            mipmap_filter: wgpu::FilterMode::Nearest,
            ..Default::default()
        });
        let layouts = Self::create_layouts(&device);
        let pipelines = Self::create_pipelines(&device, &sc_desc, &layouts, &compilation);
        let (compute_group, copy_group) = Self::bind_for_size(&device, &sampler, size, &layouts);
        let start_time = Instant::now();
        State {
            pipelines,
            layouts,
            surface,
            device,
            queue,
            sc_desc,
            swap_chain,
            size,
            sampler,
            compute_bind_group: compute_group,
            copy_bind_group: copy_group,
            start_time,
        }
    }

    fn resize(&mut self, new_size: PhysicalSize<u32>) {
        self.sc_desc.width = new_size.width;
        self.sc_desc.height = new_size.height;
        let new_swap_chain = self.device.create_swap_chain(&self.surface, &self.sc_desc);
        self.swap_chain = new_swap_chain;
        let (compute_bind_group, copy_bind_group) =
            State::bind_for_size(&self.device, &self.sampler, new_size, &self.layouts);
        self.compute_bind_group = compute_bind_group;
        self.copy_bind_group = copy_bind_group;
    }

    fn bind_for_size(
        device: &wgpu::Device,
        sampler: &wgpu::Sampler,
        size: PhysicalSize<u32>,
        layouts: &Layouts,
    ) -> (wgpu::BindGroup, wgpu::BindGroup) {
        let img = device.create_texture(&wgpu::TextureDescriptor {
            label: None,
            size: Extent3d {
                width: size.width,
                height: size.height,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8Unorm,
            usage: wgpu::TextureUsage::STORAGE | wgpu::TextureUsage::SAMPLED,
        });
        let img_view = img.create_view(&Default::default());
        let compute_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layouts.compute_bg_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: wgpu::BindingResource::TextureView(&img_view),
            }],
        });

        let copy_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layouts.copy_bg_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(&img_view),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(&sampler),
                },
            ],
        });
        (compute_bind_group, copy_bind_group)
    }

    fn create_layouts(device: &wgpu::Device) -> Layouts {
        let copy_bg_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: None,
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStage::FRAGMENT,
                    ty: wgpu::BindingType::Texture {
                        multisampled: false,
                        // Should filterable be false if we want nearest-neighbor?
                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
                        view_dimension: wgpu::TextureViewDimension::D2,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStage::FRAGMENT,
                    ty: wgpu::BindingType::Sampler {
                        filtering: false,
                        comparison: false,
                    },
                    count: None,
                },
            ],
        });
        let copy_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: None,
            bind_group_layouts: &[&copy_bg_layout],
            push_constant_ranges: &[],
        });
        let compute_bg_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Compute Bind Group Layout"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStage::COMPUTE,
                ty: wgpu::BindingType::StorageTexture {
                    access: wgpu::StorageTextureAccess::WriteOnly,
                    format: wgpu::TextureFormat::Rgba8Unorm,
                    view_dimension: wgpu::TextureViewDimension::D2,
                },
                count: None,
            }],
        });
        let compute_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: None,
                bind_group_layouts: &[&compute_bg_layout],
                push_constant_ranges: &[wgpu::PushConstantRange {
                    stages: wgpu::ShaderStage::COMPUTE,
                    range: 0..std::mem::size_of::<Config>() as u32,
                }],
            });
        Layouts {
            compute_bg_layout,
            copy_bg_layout,
            compute_pipeline_layout,
            copy_pipeline_layout,
        }
    }

    fn replace_pipelines(&mut self, compilation: &CompileResult) {
        let new_pipelines =
            State::create_pipelines(&self.device, &self.sc_desc, &self.layouts, compilation);
        self.pipelines = new_pipelines;
    }

    fn create_pipelines(
        device: &wgpu::Device,
        sc_desc: &wgpu::SwapChainDescriptor,
        layouts: &Layouts,
        compilation: &CompileResult,
    ) -> Pipelines {
        let module = State::create_shader_module(device, compilation);
        let compute_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: None,
            layout: Some(&layouts.compute_pipeline_layout),
            module: &module,
            entry_point: "main",
        });
        let copy_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: None,
            layout: Some(&layouts.copy_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &module,
                entry_point: "vs_main",
                buffers: &[],
            },
            fragment: Some(wgpu::FragmentState {
                module: &module,
                entry_point: "fs_main",
                targets: &[sc_desc.format.into()],
            }),
            primitive: wgpu::PrimitiveState::default(),
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
        });
        Pipelines {
            compute_pipeline,
            copy_pipeline,
        }
    }

    fn create_shader_module(device: &wgpu::Device, compilation: &CompileResult) -> ShaderModule {
        let shader_flags = wgpu::ShaderFlags::empty();
        let spirv_path = compilation.module.unwrap_single();
        let f = std::fs::read(spirv_path).expect("spirv should exist");
        let shader = make_spirv(&f);
        device.create_shader_module(&wgpu::ShaderModuleDescriptor {
            label: None,
            source: shader,
            flags: shader_flags,
        })
    }

    fn render_frame(&self) {
        let frame = self
            .swap_chain
            .get_current_frame()
            .expect("error getting frame from swap chain")
            .output;

        let i_time: f32 = 0.5 + self.start_time.elapsed().as_micros() as f32 * 1e-6;
        let size = self.size;
        let config = Config {
            width: size.width,
            height: size.height,
            time: i_time,
        };

        let mut encoder = self.device.create_command_encoder(&Default::default());
        {
            let mut cpass = encoder.begin_compute_pass(&Default::default());
            cpass.set_pipeline(&self.pipelines.compute_pipeline);
            cpass.set_push_constants(0, bytemuck::bytes_of(&config));
            cpass.set_bind_group(0, &self.compute_bind_group, &[]);
            cpass.dispatch(size.width / 16, size.height / 16, 1);
        }
        {
            let mut rpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: None,
                color_attachments: &[wgpu::RenderPassColorAttachment {
                    view: &frame.view,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::GREEN),
                        store: true,
                    },
                }],
                depth_stencil_attachment: None,
            });
            rpass.set_pipeline(&self.pipelines.copy_pipeline);
            rpass.set_bind_group(0, &self.copy_bind_group, &[]);
            rpass.draw(0..3, 0..2);
        }
        self.queue.submit(Some(encoder.finish()));
    }
}

async fn run(
    event_loop: EventLoop<CompileResult>,
    window: Window,
    initial_compilation: CompileResult,
) {
    let mut state = State::new(&window, initial_compilation).await;

    event_loop.run(move |event, _, control_flow| {
        // TODO: this may be excessive polling. It really should be synchronized with
        // swapchain presentation, but that's currently underbaked in wgpu.
        *control_flow = ControlFlow::Poll;
        match event {
            Event::RedrawRequested(_) => state.render_frame(),
            Event::MainEventsCleared => {
                window.request_redraw();
            }
            Event::WindowEvent {
                event: WindowEvent::CloseRequested,
                ..
            } => *control_flow = ControlFlow::Exit,
            Event::UserEvent(result) => state.replace_pipelines(&result),
            Event::WindowEvent { window_id, event } if window_id == window.id() => match event {
                WindowEvent::Resized(new_size) => {
                    state.resize(new_size);
                    window.request_redraw();
                }
                _ => (),
            },
            _ => (),
        }
    });
}

fn main() {
    let (initial_tx, initial_rx) = std::sync::mpsc::sync_channel(0);
    let event_loop = EventLoop::with_user_event();
    let proxy = event_loop.create_proxy();
    // This is going to be cleaned up significantly by https://github.com/EmbarkStudios/rust-gpu/pull/663
    let has_sent_first = AtomicBool::new(false);
    // Watch for changes on a background thread
    let thread = std::thread::spawn(|| {
        SpirvBuilder::new("./shaders", "spirv-unknown-vulkan1.2")
            .print_metadata(spirv_builder::MetadataPrintout::None)
            .capability(spirv_builder::Capability::StorageImageWriteWithoutFormat)
            .watch(move |result| {
                if let Ok(_) =
                    // TODO: This use of atomics is signic
                    has_sent_first.compare_exchange(
                        false,
                        true,
                        Ordering::SeqCst,
                        Ordering::SeqCst,
                    )
                {
                    initial_tx.send(result).unwrap();
                } else {
                    proxy
                        .send_event(result)
                        .expect("Event loop should still be running");
                }
            })
            .expect("Correctly setup")
    });
    std::mem::forget(thread);
    let initial = initial_rx
        .recv()
        .expect("Watching should get an initial shader");
    let window = Window::new(&event_loop).unwrap();
    pollster::block_on(run(event_loop, window, initial));
}
<commit_msg>Resize upon resize<commit_after>// Copyright 2021 Google LLC
//
// 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
//
//     https://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.
//
// Also licensed under MIT license, at your choice.

//! A simple compute shader example that draws into a window, based on wgpu.

use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Instant;

use rust_gpu_toy_shared::Config;
use wgpu::util::make_spirv;
use wgpu::{Extent3d, ShaderModule};

use winit::dpi::PhysicalSize;
use winit::{
    event::{Event, WindowEvent},
    event_loop::{ControlFlow, EventLoop},
    window::Window,
};

use spirv_builder::{CompileResult, SpirvBuilder};

struct Pipelines {
    compute_pipeline: wgpu::ComputePipeline,

    copy_pipeline: wgpu::RenderPipeline,
}

struct Layouts {
    compute_bg_layout: wgpu::BindGroupLayout,
    copy_bg_layout: wgpu::BindGroupLayout,

    compute_pipeline_layout: wgpu::PipelineLayout,
    copy_pipeline_layout: wgpu::PipelineLayout,
}

/// State of the runner, inspired by
/// https://sotrh.github.io/learn-wgpu/
struct State {
    pipelines: Pipelines,
    layouts: Layouts,

    surface: wgpu::Surface,
    device: wgpu::Device,
    queue: wgpu::Queue,

    sc_desc: wgpu::SwapChainDescriptor,
    swap_chain: wgpu::SwapChain,
    size: winit::dpi::PhysicalSize<u32>,

    sampler: wgpu::Sampler,

    compute_bind_group: wgpu::BindGroup,
    copy_bind_group: wgpu::BindGroup,

    start_time: Instant,
}

impl State {
    async fn new(window: &Window, compilation: CompileResult) -> Self {
        let size = window.inner_size();
        let instance = wgpu::Instance::new(wgpu::BackendBit::PRIMARY);
        let surface = unsafe { instance.create_surface(window) };
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: Default::default(),
                compatible_surface: Some(&surface),
            })
            .await
            .expect("error finding adapter");
        let features = wgpu::Features::PUSH_CONSTANTS;
        let limits = wgpu::Limits {
            max_push_constant_size: std::mem::size_of::<Config>() as u32,
            ..Default::default()
        };
        let (device, queue) = adapter
            .request_device(
                &wgpu::DeviceDescriptor {
                    label: None,
                    features,
                    limits,
                },
                None,
            )
            .await
            .expect("error creating device");
        let swapchain_format = adapter.get_swap_chain_preferred_format(&surface).unwrap();
        let sc_desc = wgpu::SwapChainDescriptor {
            usage: wgpu::TextureUsage::RENDER_ATTACHMENT,
            format: swapchain_format,
            width: size.width,
            height: size.height,
            present_mode: wgpu::PresentMode::Mailbox,
        };

        let swap_chain = device.create_swap_chain(&surface, &sc_desc);

        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            address_mode_u: wgpu::AddressMode::ClampToEdge,
            address_mode_v: wgpu::AddressMode::ClampToEdge,
            address_mode_w: wgpu::AddressMode::ClampToEdge,
            mag_filter: wgpu::FilterMode::Nearest,
            min_filter: wgpu::FilterMode::Nearest,
            mipmap_filter: wgpu::FilterMode::Nearest,
            ..Default::default()
        });
        let layouts = Self::create_layouts(&device);
        let pipelines = Self::create_pipelines(&device, &sc_desc, &layouts, &compilation);
        let (compute_group, copy_group) = Self::bind_for_size(&device, &sampler, size, &layouts);
        let start_time = Instant::now();
        State {
            pipelines,
            layouts,
            surface,
            device,
            queue,
            sc_desc,
            swap_chain,
            size,
            sampler,
            compute_bind_group: compute_group,
            copy_bind_group: copy_group,
            start_time,
        }
    }

    fn resize(&mut self, new_size: PhysicalSize<u32>) {
        self.sc_desc.width = new_size.width;
        self.sc_desc.height = new_size.height;
        self.size = new_size;
        let new_swap_chain = self.device.create_swap_chain(&self.surface, &self.sc_desc);
        self.swap_chain = new_swap_chain;
        let (compute_bind_group, copy_bind_group) =
            State::bind_for_size(&self.device, &self.sampler, new_size, &self.layouts);
        self.compute_bind_group = compute_bind_group;
        self.copy_bind_group = copy_bind_group;
    }

    fn bind_for_size(
        device: &wgpu::Device,
        sampler: &wgpu::Sampler,
        size: PhysicalSize<u32>,
        layouts: &Layouts,
    ) -> (wgpu::BindGroup, wgpu::BindGroup) {
        let img = device.create_texture(&wgpu::TextureDescriptor {
            label: None,
            size: Extent3d {
                width: size.width,
                height: size.height,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8Unorm,
            usage: wgpu::TextureUsage::STORAGE | wgpu::TextureUsage::SAMPLED,
        });
        let img_view = img.create_view(&Default::default());
        let compute_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layouts.compute_bg_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: wgpu::BindingResource::TextureView(&img_view),
            }],
        });

        let copy_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layouts.copy_bg_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(&img_view),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(&sampler),
                },
            ],
        });
        (compute_bind_group, copy_bind_group)
    }

    fn create_layouts(device: &wgpu::Device) -> Layouts {
        let copy_bg_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: None,
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStage::FRAGMENT,
                    ty: wgpu::BindingType::Texture {
                        multisampled: false,
                        // Should filterable be false if we want nearest-neighbor?
                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
                        view_dimension: wgpu::TextureViewDimension::D2,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStage::FRAGMENT,
                    ty: wgpu::BindingType::Sampler {
                        filtering: false,
                        comparison: false,
                    },
                    count: None,
                },
            ],
        });
        let copy_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: None,
            bind_group_layouts: &[&copy_bg_layout],
            push_constant_ranges: &[],
        });
        let compute_bg_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Compute Bind Group Layout"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStage::COMPUTE,
                ty: wgpu::BindingType::StorageTexture {
                    access: wgpu::StorageTextureAccess::WriteOnly,
                    format: wgpu::TextureFormat::Rgba8Unorm,
                    view_dimension: wgpu::TextureViewDimension::D2,
                },
                count: None,
            }],
        });
        let compute_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: None,
                bind_group_layouts: &[&compute_bg_layout],
                push_constant_ranges: &[wgpu::PushConstantRange {
                    stages: wgpu::ShaderStage::COMPUTE,
                    range: 0..std::mem::size_of::<Config>() as u32,
                }],
            });
        Layouts {
            compute_bg_layout,
            copy_bg_layout,
            compute_pipeline_layout,
            copy_pipeline_layout,
        }
    }

    fn replace_pipelines(&mut self, compilation: &CompileResult) {
        let new_pipelines =
            State::create_pipelines(&self.device, &self.sc_desc, &self.layouts, compilation);
        self.pipelines = new_pipelines;
    }

    fn create_pipelines(
        device: &wgpu::Device,
        sc_desc: &wgpu::SwapChainDescriptor,
        layouts: &Layouts,
        compilation: &CompileResult,
    ) -> Pipelines {
        let module = State::create_shader_module(device, compilation);
        let compute_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: None,
            layout: Some(&layouts.compute_pipeline_layout),
            module: &module,
            entry_point: "main",
        });
        let copy_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: None,
            layout: Some(&layouts.copy_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &module,
                entry_point: "vs_main",
                buffers: &[],
            },
            fragment: Some(wgpu::FragmentState {
                module: &module,
                entry_point: "fs_main",
                targets: &[sc_desc.format.into()],
            }),
            primitive: wgpu::PrimitiveState::default(),
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
        });
        Pipelines {
            compute_pipeline,
            copy_pipeline,
        }
    }

    fn create_shader_module(device: &wgpu::Device, compilation: &CompileResult) -> ShaderModule {
        let shader_flags = wgpu::ShaderFlags::empty();
        let spirv_path = compilation.module.unwrap_single();
        let f = std::fs::read(spirv_path).expect("spirv should exist");
        let shader = make_spirv(&f);
        device.create_shader_module(&wgpu::ShaderModuleDescriptor {
            label: None,
            source: shader,
            flags: shader_flags,
        })
    }

    fn render_frame(&self) {
        let frame = self
            .swap_chain
            .get_current_frame()
            .expect("error getting frame from swap chain")
            .output;

        let i_time: f32 = 0.5 + self.start_time.elapsed().as_micros() as f32 * 1e-6;
        let size = self.size;
        let config = Config {
            width: size.width,
            height: size.height,
            time: i_time,
        };

        let mut encoder = self.device.create_command_encoder(&Default::default());
        {
            let mut cpass = encoder.begin_compute_pass(&Default::default());
            cpass.set_pipeline(&self.pipelines.compute_pipeline);
            cpass.set_push_constants(0, bytemuck::bytes_of(&config));
            cpass.set_bind_group(0, &self.compute_bind_group, &[]);
            cpass.dispatch(size.width / 16, size.height / 16, 1);
        }
        {
            let mut rpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: None,
                color_attachments: &[wgpu::RenderPassColorAttachment {
                    view: &frame.view,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::GREEN),
                        store: true,
                    },
                }],
                depth_stencil_attachment: None,
            });
            rpass.set_pipeline(&self.pipelines.copy_pipeline);
            rpass.set_bind_group(0, &self.copy_bind_group, &[]);
            rpass.draw(0..3, 0..2);
        }
        self.queue.submit(Some(encoder.finish()));
    }
}

async fn run(
    event_loop: EventLoop<CompileResult>,
    window: Window,
    initial_compilation: CompileResult,
) {
    let mut state = State::new(&window, initial_compilation).await;

    event_loop.run(move |event, _, control_flow| {
        // TODO: this may be excessive polling. It really should be synchronized with
        // swapchain presentation, but that's currently underbaked in wgpu.
        *control_flow = ControlFlow::Poll;
        match event {
            Event::RedrawRequested(_) => state.render_frame(),
            Event::MainEventsCleared => {
                window.request_redraw();
            }
            Event::WindowEvent {
                event: WindowEvent::CloseRequested,
                ..
            } => *control_flow = ControlFlow::Exit,
            Event::UserEvent(result) => state.replace_pipelines(&result),
            Event::WindowEvent { window_id, event } if window_id == window.id() => match event {
                WindowEvent::Resized(new_size) => {
                    state.resize(new_size);
                    window.request_redraw();
                }
                _ => (),
            },
            _ => (),
        }
    });
}

fn main() {
    let (initial_tx, initial_rx) = std::sync::mpsc::sync_channel(0);
    let event_loop = EventLoop::with_user_event();
    let proxy = event_loop.create_proxy();
    // This is going to be cleaned up significantly by https://github.com/EmbarkStudios/rust-gpu/pull/663
    let has_sent_first = AtomicBool::new(false);
    // Watch for changes on a background thread
    let thread = std::thread::spawn(|| {
        SpirvBuilder::new("./shaders", "spirv-unknown-vulkan1.2")
            .print_metadata(spirv_builder::MetadataPrintout::None)
            .capability(spirv_builder::Capability::StorageImageWriteWithoutFormat)
            .watch(move |result| {
                if let Ok(_) =
                    // TODO: This use of atomics is signic
                    has_sent_first.compare_exchange(
                        false,
                        true,
                        Ordering::SeqCst,
                        Ordering::SeqCst,
                    )
                {
                    initial_tx.send(result).unwrap();
                } else {
                    proxy
                        .send_event(result)
                        .expect("Event loop should still be running");
                }
            })
            .expect("Correctly setup")
    });
    std::mem::forget(thread);
    let initial = initial_rx
        .recv()
        .expect("Watching should get an initial shader");
    let window = Window::new(&event_loop).unwrap();
    pollster::block_on(run(event_loop, window, initial));
}
<|endoftext|>


<commit_before>// Copyright 2013 The rust-pcre authors.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern mod extra;

use extra::enum_set::{EnumSet};
use std::c_str::{CString};
use std::libc::{c_int, c_char, c_void, c_uchar};
use std::ptr;
use std::result::{Result};

mod native;

pub type compile_options = c_int;
pub type exec_options = c_int;
pub type fullinfo_field = c_int;
pub struct pcre;
pub type pcre_error = c_int;
pub struct pcre_extra;
pub type study_options = c_int;

pub static PCRE_UTF8: c_int = 0x00000800;
pub static PCRE_NO_UTF8_CHECK: c_int = 0x00002000;

pub static PCRE_ERROR_NOMATCH: pcre_error = -1;
pub static PCRE_ERROR_NULL: pcre_error = -2;

pub static PCRE_INFO_OPTIONS: fullinfo_field = 0;
pub static PCRE_INFO_SIZE: fullinfo_field = 1;
pub static PCRE_INFO_CAPTURECOUNT: fullinfo_field = 2;
pub static PCRE_INFO_BACKREFMAX: fullinfo_field = 3;
pub static PCRE_INFO_FIRSTBYTE: fullinfo_field = 4;
pub static PCRE_INFO_FIRSTCHAR: fullinfo_field = 4;
pub static PCRE_INFO_FIRSTTABLE: fullinfo_field = 5;
pub static PCRE_INFO_LASTLITERAL: fullinfo_field = 6;
pub static PCRE_INFO_NAMEENTRYSIZE: fullinfo_field = 7;
pub static PCRE_INFO_NAMECOUNT: fullinfo_field = 8;
pub static PCRE_INFO_NAMETABLE: fullinfo_field = 9;
pub static PCRE_INFO_STUDYSIZE: fullinfo_field = 10;
pub static PCRE_INFO_DEFAULT_TABLES: fullinfo_field = 11;
pub static PCRE_INFO_OKPARTIAL: fullinfo_field = 12;
pub static PCRE_INFO_JCHANGED: fullinfo_field = 13;
pub static PCRE_INFO_HASCRORLF: fullinfo_field = 14;
pub static PCRE_INFO_MINLENGTH: fullinfo_field = 15;
pub static PCRE_INFO_JIT: fullinfo_field = 16;
pub static PCRE_INFO_JITSIZE: fullinfo_field = 17;
pub static PCRE_INFO_MAXLOOKBEHIND: fullinfo_field = 18;
pub static PCRE_INFO_FIRSTCHARACTER: fullinfo_field = 19;
pub static PCRE_INFO_FIRSTCHARACTERFLAGS: fullinfo_field = 20;
pub static PCRE_INFO_REQUIREDCHAR: fullinfo_field = 21;
pub static PCRE_INFO_REQUIREDCHARFLAGS: fullinfo_field = 22;
pub static PCRE_INFO_MATCHLIMIT: fullinfo_field = 23;
pub static PCRE_INFO_RECURSIONLIMIT: fullinfo_field = 24;

pub unsafe fn pcre_compile(pattern: *c_char, options: &EnumSet<::CompileOption>, tableptr: *c_uchar) -> Result<*mut pcre, (Option<~str>, c_int)> {
    assert!(ptr::is_not_null(pattern));
    let converted_options = options.iter().fold(0, |converted_options, option| converted_options | (option as compile_options)) | PCRE_UTF8 | PCRE_NO_UTF8_CHECK;
    let mut err: *c_char = ptr::null();
    let mut erroffset: c_int = 0;
    let code = native::pcre_compile(pattern, converted_options, &mut err, &mut erroffset, tableptr);

    if ptr::is_null(code) {
        // "Otherwise, if  compilation  of  a  pattern fails, pcre_compile() returns
        // NULL, and sets the variable pointed to by errptr to point to a textual
        // error message. This is a static string that is part of the library. You
        // must not try to free it."
        // http://pcre.org/pcre.txt
        let err_cstring = CString::new(err, false);
        //let err_str = match err_cstring.as_str() {
        //    None          => format!("compilation failed at offset {:u}", erroffset as uint),
        //    Some(err_str) => format!("compilation failed at offset {:u}: {:s}", erroffset as uint, err_str)
        //};

        match err_cstring.as_str() {
            None => Err((None, erroffset)),
            Some(err_str) => Err((Some(err_str.to_owned()), erroffset))
        }
    } else {
        assert!(ptr::is_not_null(code));
        assert_eq!(erroffset, 0);

        Ok(code)
    }
}

pub unsafe fn pcre_exec(code: *pcre, extra: *pcre_extra, subject: *c_char, length: c_int, startoffset: c_int, options: &EnumSet<::ExecOption>, ovector: *mut c_int, ovecsize: c_int) -> c_int {
    assert!(ptr::is_not_null(code));
    assert!(ovecsize >= 0 && ovecsize % 3 == 0);
    let converted_options = options.iter().fold(0, |converted_options, option| converted_options | (option as compile_options)) | PCRE_NO_UTF8_CHECK;
    let rc = native::pcre_exec(code, extra, subject, length, startoffset, converted_options, ovector, ovecsize);
    if rc == PCRE_ERROR_NOMATCH {
        return -1;
    } else if rc < 0 && rc != PCRE_ERROR_NULL {
        fail!("pcre_exec");
    }

    rc
}

pub unsafe fn pcre_free(ptr: *mut c_void) {
    native::pcre_free(ptr);
}

pub unsafe fn pcre_free_study(extra: *mut pcre_extra) {
    native::pcre_free_study(extra);
}

pub unsafe fn pcre_fullinfo(code: *pcre, extra: *pcre_extra, what: fullinfo_field, where: *mut c_void) {
    assert!(ptr::is_not_null(code));
    let rc = native::pcre_fullinfo(code, extra, what, where);
    if rc < 0 && rc != PCRE_ERROR_NULL {
        fail!("pcre_fullinfo");
    }
}

pub unsafe fn pcre_refcount(code: *mut ::detail::pcre, adjust: c_int) -> c_int {
    assert!(ptr::is_not_null(code));
    let curr_refcount = native::pcre_refcount(code, 0);
    if curr_refcount + adjust < 0 {
        fail!("refcount underflow");
    } else if curr_refcount + adjust > 65535 {
        fail!("refcount overflow");
    }
    native::pcre_refcount(code, adjust)
}

pub unsafe fn pcre_study(code: *::detail::pcre, options: &EnumSet<::StudyOption>) -> *mut ::detail::pcre_extra {
    assert!(ptr::is_not_null(code));
    let converted_options = options.iter().fold(0, |converted_options, option| converted_options | (option as study_options));
    let mut err: *c_char = ptr::null();
    let extra = native::pcre_study(code, converted_options, &mut err);
    // "The third argument for pcre_study() is a pointer for an error message. If
    // studying succeeds (even if no data is returned), the variable it points to is
    // set to NULL. Otherwise it is set to point to a textual error message. This is
    // a static string that is part of the library. You must not try to free it."
    // http://pcre.org/pcre.txt
    if ptr::is_not_null(err) {
        let err_cstring = CString::new(err, false);
        match err_cstring.as_str() {
            None          => error!("pcre_study() failed"),
            Some(err_str) => error!("pcre_study() failed: {}", err_str)
        }
        fail!("pcre_study");
    }
    assert!(ptr::is_null(err));

    extra
}

pub fn pcre_version() -> ~str {
    let version_cstring = unsafe { CString::new(native::pcre_version(), false) };
    version_cstring.as_str().unwrap().to_owned()
}
<commit_msg>Remove unused, commented-out code<commit_after>// Copyright 2013 The rust-pcre authors.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern mod extra;

use extra::enum_set::{EnumSet};
use std::c_str::{CString};
use std::libc::{c_int, c_char, c_void, c_uchar};
use std::ptr;
use std::result::{Result};

mod native;

pub type compile_options = c_int;
pub type exec_options = c_int;
pub type fullinfo_field = c_int;
pub struct pcre;
pub type pcre_error = c_int;
pub struct pcre_extra;
pub type study_options = c_int;

pub static PCRE_UTF8: c_int = 0x00000800;
pub static PCRE_NO_UTF8_CHECK: c_int = 0x00002000;

pub static PCRE_ERROR_NOMATCH: pcre_error = -1;
pub static PCRE_ERROR_NULL: pcre_error = -2;

pub static PCRE_INFO_OPTIONS: fullinfo_field = 0;
pub static PCRE_INFO_SIZE: fullinfo_field = 1;
pub static PCRE_INFO_CAPTURECOUNT: fullinfo_field = 2;
pub static PCRE_INFO_BACKREFMAX: fullinfo_field = 3;
pub static PCRE_INFO_FIRSTBYTE: fullinfo_field = 4;
pub static PCRE_INFO_FIRSTCHAR: fullinfo_field = 4;
pub static PCRE_INFO_FIRSTTABLE: fullinfo_field = 5;
pub static PCRE_INFO_LASTLITERAL: fullinfo_field = 6;
pub static PCRE_INFO_NAMEENTRYSIZE: fullinfo_field = 7;
pub static PCRE_INFO_NAMECOUNT: fullinfo_field = 8;
pub static PCRE_INFO_NAMETABLE: fullinfo_field = 9;
pub static PCRE_INFO_STUDYSIZE: fullinfo_field = 10;
pub static PCRE_INFO_DEFAULT_TABLES: fullinfo_field = 11;
pub static PCRE_INFO_OKPARTIAL: fullinfo_field = 12;
pub static PCRE_INFO_JCHANGED: fullinfo_field = 13;
pub static PCRE_INFO_HASCRORLF: fullinfo_field = 14;
pub static PCRE_INFO_MINLENGTH: fullinfo_field = 15;
pub static PCRE_INFO_JIT: fullinfo_field = 16;
pub static PCRE_INFO_JITSIZE: fullinfo_field = 17;
pub static PCRE_INFO_MAXLOOKBEHIND: fullinfo_field = 18;
pub static PCRE_INFO_FIRSTCHARACTER: fullinfo_field = 19;
pub static PCRE_INFO_FIRSTCHARACTERFLAGS: fullinfo_field = 20;
pub static PCRE_INFO_REQUIREDCHAR: fullinfo_field = 21;
pub static PCRE_INFO_REQUIREDCHARFLAGS: fullinfo_field = 22;
pub static PCRE_INFO_MATCHLIMIT: fullinfo_field = 23;
pub static PCRE_INFO_RECURSIONLIMIT: fullinfo_field = 24;

pub unsafe fn pcre_compile(pattern: *c_char, options: &EnumSet<::CompileOption>, tableptr: *c_uchar) -> Result<*mut pcre, (Option<~str>, c_int)> {
    assert!(ptr::is_not_null(pattern));
    let converted_options = options.iter().fold(0, |converted_options, option| converted_options | (option as compile_options)) | PCRE_UTF8 | PCRE_NO_UTF8_CHECK;
    let mut err: *c_char = ptr::null();
    let mut erroffset: c_int = 0;
    let code = native::pcre_compile(pattern, converted_options, &mut err, &mut erroffset, tableptr);

    if ptr::is_null(code) {
        // "Otherwise, if  compilation  of  a  pattern fails, pcre_compile() returns
        // NULL, and sets the variable pointed to by errptr to point to a textual
        // error message. This is a static string that is part of the library. You
        // must not try to free it."
        // http://pcre.org/pcre.txt
        let err_cstring = CString::new(err, false);

        match err_cstring.as_str() {
            None => Err((None, erroffset)),
            Some(err_str) => Err((Some(err_str.to_owned()), erroffset))
        }
    } else {
        assert!(ptr::is_not_null(code));
        assert_eq!(erroffset, 0);

        Ok(code)
    }
}

pub unsafe fn pcre_exec(code: *pcre, extra: *pcre_extra, subject: *c_char, length: c_int, startoffset: c_int, options: &EnumSet<::ExecOption>, ovector: *mut c_int, ovecsize: c_int) -> c_int {
    assert!(ptr::is_not_null(code));
    assert!(ovecsize >= 0 && ovecsize % 3 == 0);
    let converted_options = options.iter().fold(0, |converted_options, option| converted_options | (option as compile_options)) | PCRE_NO_UTF8_CHECK;
    let rc = native::pcre_exec(code, extra, subject, length, startoffset, converted_options, ovector, ovecsize);
    if rc == PCRE_ERROR_NOMATCH {
        return -1;
    } else if rc < 0 && rc != PCRE_ERROR_NULL {
        fail!("pcre_exec");
    }

    rc
}

pub unsafe fn pcre_free(ptr: *mut c_void) {
    native::pcre_free(ptr);
}

pub unsafe fn pcre_free_study(extra: *mut pcre_extra) {
    native::pcre_free_study(extra);
}

pub unsafe fn pcre_fullinfo(code: *pcre, extra: *pcre_extra, what: fullinfo_field, where: *mut c_void) {
    assert!(ptr::is_not_null(code));
    let rc = native::pcre_fullinfo(code, extra, what, where);
    if rc < 0 && rc != PCRE_ERROR_NULL {
        fail!("pcre_fullinfo");
    }
}

pub unsafe fn pcre_refcount(code: *mut ::detail::pcre, adjust: c_int) -> c_int {
    assert!(ptr::is_not_null(code));
    let curr_refcount = native::pcre_refcount(code, 0);
    if curr_refcount + adjust < 0 {
        fail!("refcount underflow");
    } else if curr_refcount + adjust > 65535 {
        fail!("refcount overflow");
    }
    native::pcre_refcount(code, adjust)
}

pub unsafe fn pcre_study(code: *::detail::pcre, options: &EnumSet<::StudyOption>) -> *mut ::detail::pcre_extra {
    assert!(ptr::is_not_null(code));
    let converted_options = options.iter().fold(0, |converted_options, option| converted_options | (option as study_options));
    let mut err: *c_char = ptr::null();
    let extra = native::pcre_study(code, converted_options, &mut err);
    // "The third argument for pcre_study() is a pointer for an error message. If
    // studying succeeds (even if no data is returned), the variable it points to is
    // set to NULL. Otherwise it is set to point to a textual error message. This is
    // a static string that is part of the library. You must not try to free it."
    // http://pcre.org/pcre.txt
    if ptr::is_not_null(err) {
        let err_cstring = CString::new(err, false);
        match err_cstring.as_str() {
            None          => error!("pcre_study() failed"),
            Some(err_str) => error!("pcre_study() failed: {}", err_str)
        }
        fail!("pcre_study");
    }
    assert!(ptr::is_null(err));

    extra
}

pub fn pcre_version() -> ~str {
    let version_cstring = unsafe { CString::new(native::pcre_version(), false) };
    version_cstring.as_str().unwrap().to_owned()
}
<|endoftext|>


<commit_before>pub static VSHADER_SOURCE: &'static str = r#"
    #version 150
    in vec2 position0;
    in vec2 position1;
    in vec2 position2;
    in vec4 inner_color0;
    in vec4 inner_color1;
    in float falloff0;
    in float falloff1;
    in vec4 falloff_color0;
    in vec4 falloff_color1;
    in float falloff_radius0;
    in float falloff_radius1;
    in float inner_radius0;
    in float inner_radius1;
    in int accuracy;

    out vec2 gposition0;
    out vec2 gposition1;
    out vec2 gposition2;
    out vec4 ginner_color0;
    out vec4 ginner_color1;
    out float gfalloff0;
    out float gfalloff1;
    out vec4 gfalloff_color0;
    out vec4 gfalloff_color1;
    out float gfalloff_radius0;
    out float gfalloff_radius1;
    out float ginner_radius0;
    out float ginner_radius1;
    out int gaccuracy;

    void main() {
        gposition0 = position0;
        gposition1 = position1;
        gposition2 = position2;
        ginner_color0 = inner_color0;
        ginner_color1 = inner_color1;
        gfalloff0 = falloff0;
        gfalloff1 = falloff1;
        gfalloff_color0 = falloff_color0;
        gfalloff_color1 = falloff_color1;
        gfalloff_radius0 = falloff_radius0;
        gfalloff_radius1 = falloff_radius1;
        ginner_radius0 = inner_radius0;
        ginner_radius1 = inner_radius1;
        gaccuracy = accuracy;
        gl_Position = vec4(0.0, 0.0, 0.0, 1.0);
    }
"#;

pub static GSHADER_SOURCE: &'static str = r#"
    #version 150

    layout(points) in;
    layout(triangle_strip, max_vertices = 9) out;

    in vec2 gposition0[1];
    in vec2 gposition1[1];
    in vec2 gposition2[1];
    in vec4 ginner_color0[1];
    in vec4 ginner_color1[1];
    in float gfalloff0[1];
    in float gfalloff1[1];
    in vec4 gfalloff_color0[1];
    in vec4 gfalloff_color1[1];
    in float gfalloff_radius0[1];
    in float gfalloff_radius1[1];
    in float ginner_radius0[1];
    in float ginner_radius1[1];
    in int gaccuracy[1];

    out vec2 fposition0;
    out vec2 fposition1;
    out vec2 fposition2;
    out vec4 finner_color0;
    out vec4 finner_color1;
    out float ffalloff0;
    out float ffalloff1;
    out vec4 ffalloff_color0;
    out vec4 ffalloff_color1;
    out float ffalloff_radius0;
    out float ffalloff_radius1;
    out float finner_radius0;
    out float finner_radius1;
    out vec2 realpos;
    flat out int faccuracy;

    void main() {
        finner_color0 = ginner_color0[0];
        finner_color1 = ginner_color1[0];
        ffalloff0 = gfalloff0[0];
        ffalloff1 = gfalloff1[0];
        ffalloff_color0 = gfalloff_color0[0];
        ffalloff_color1 = gfalloff_color1[0];
        ffalloff_radius0 = gfalloff_radius0[0];
        ffalloff_radius1 = gfalloff_radius1[0];
        finner_radius0 = ginner_radius0[0];
        finner_radius1 = ginner_radius1[0];
        faccuracy = gaccuracy[0];

        vec2 l0 = normalize(gposition1[0] - gposition0[0]);
        vec2 l1 = normalize(gposition2[0] - gposition1[0]);
        vec2 l2 = normalize(gposition0[0] - gposition2[0]);

        vec2 b0 = normalize(l0 - l2);
        vec2 b1 = normalize(l0 - l1);
        vec2 b2 = normalize(l2 - l1);

        float radius0 = finner_radius0 + ffalloff_radius0;
        float radius2 = finner_radius1 + ffalloff_radius1;
        float radius1 = (radius0 + radius2) * 0.5;

        vec4 e0 = vec4(gposition0[0] + radius0 * vec2(b0.y, -b0.x), 0.0, 1.0);
        vec4 e1 = vec4(gposition0[0] + radius0 * vec2(-b0.y, b0.x), 0.0, 1.0);
        vec4 e2 = vec4(gposition1[0] + radius1 * b1, 0.0, 1.0);
        vec4 e3 = vec4(gposition2[0] + radius2 * vec2(b2.y, -b2.x), 0.0, 1.0);
        vec4 e4 = vec4(gposition2[0] + radius2 * vec2(-b2.y, b2.x), 0.0, 1.0);

        gl_Position = e2;
        realpos = e2.xy;
        EmitVertex();

        gl_Position = e1;
        realpos = e1.xy;
        EmitVertex();

        gl_Position = e0;
        realpos = e0.xy;
        EmitVertex();

        EndPrimitive();

        gl_Position = e2;
        realpos = e2.xy;
        EmitVertex();

        gl_Position = e0;
        realpos = e0.xy;
        EmitVertex();

        gl_Position = e4;
        realpos = e4.xy;
        EmitVertex();

        EndPrimitive();

        gl_Position = e2;
        realpos = e2.xy;
        EmitVertex();

        gl_Position = e4;
        realpos = e4.xy;
        EmitVertex();

        gl_Position = e3;
        realpos = e3.xy;
        EmitVertex();

        EndPrimitive();
    }
"#;

pub static FSHADER_SOURCE: &'static str = r#"
    #version 150

    in vec2 fposition0;
    in vec2 fposition1;
    in vec2 fposition2;
    in vec4 finner_color0;
    in vec4 finner_color1;
    in float ffalloff0;
    in float ffalloff1;
    in vec4 ffalloff_color0;
    in vec4 ffalloff_color1;
    in float ffalloff_radius0;
    in float ffalloff_radius1;
    in float finner_radius0;
    in float finner_radius1;
    in vec2 realpos;
    flat in int faccuracy;

    out vec4 color;

    void main() {
        bool best_begin;
        float best_distance;
        float best_begin_t = 0.0;
        float best_end_t = 1.0;

        float beginning_distance = length(fposition0 - realpos);
        float ending_distance = length(fposition2 - realpos);
        if (beginning_distance < ending_distance) {
            best_begin = true;
            best_distance = beginning_distance;
        } else {
            best_begin = false;
            best_distance = ending_distance;
        }

        for (int i = 0; i < faccuracy; i++) {
            // Solve the bezier curve distance between beginning and end.
            float middle_t = (best_begin_t + best_end_t) * 0.5;
            float distance = length(
                pow(1.0 - middle_t, 2) * fposition0 +
                2.0 * middle_t * (1.0 - middle_t) * fposition1 +
                pow(middle_t, 2) * fposition2 -
                realpos);

            if (distance < best_distance) {
                if (best_begin) {
                    best_begin = false;
                    best_distance = distance;
                    best_end_t = middle_t;
                } else {
                    best_begin = true;
                    best_distance = distance;
                    best_begin_t = middle_t;
                }
            } else {
                if (best_begin) {
                    best_end_t = middle_t;
                } else {
                    best_begin_t = middle_t;
                }
            }
        }

        float best_t;
        if (best_begin) {
            best_t = best_begin_t;
        } else {
            best_t = best_end_t;
        }

        vec4 inner_color = finner_color0 + best_t * (finner_color1 - finner_color0);
        float falloff = ffalloff0 + best_t * (ffalloff1 - ffalloff0);
        vec4 falloff_color = ffalloff_color0 + best_t * (ffalloff_color1 - ffalloff_color0);
        float falloff_radius = ffalloff_radius0 + best_t * (ffalloff_radius1 - ffalloff_radius0);
        float inner_radius = finner_radius0 + best_t * (finner_radius1 - finner_radius0);

        if (best_distance <= inner_radius) {
            float travel = best_distance / inner_radius;
            // Manually interpolate the inner color into the falloff color.
            color = inner_color * (1.0 - travel) + falloff_color * travel;
        } else {
            color = vec4(falloff_color.xyz,
                falloff_color.a * max(0.0, 1.0 - pow((best_distance - inner_radius) / falloff_radius, falloff)));
        }
    }
"#;
<commit_msg>shaders largely working now except for a minor segmentation bug<commit_after>pub static VSHADER_SOURCE: &'static str = r#"
    #version 150
    in vec2 position0;
    in vec2 position1;
    in vec2 position2;
    in vec4 inner_color0;
    in vec4 inner_color1;
    in float falloff0;
    in float falloff1;
    in vec4 falloff_color0;
    in vec4 falloff_color1;
    in float falloff_radius0;
    in float falloff_radius1;
    in float inner_radius0;
    in float inner_radius1;
    in int accuracy;

    out vec2 gposition0;
    out vec2 gposition1;
    out vec2 gposition2;
    out vec4 ginner_color0;
    out vec4 ginner_color1;
    out float gfalloff0;
    out float gfalloff1;
    out vec4 gfalloff_color0;
    out vec4 gfalloff_color1;
    out float gfalloff_radius0;
    out float gfalloff_radius1;
    out float ginner_radius0;
    out float ginner_radius1;
    out int gaccuracy;

    void main() {
        gposition0 = position0;
        gposition1 = position1;
        gposition2 = position2;
        ginner_color0 = inner_color0;
        ginner_color1 = inner_color1;
        gfalloff0 = falloff0;
        gfalloff1 = falloff1;
        gfalloff_color0 = falloff_color0;
        gfalloff_color1 = falloff_color1;
        gfalloff_radius0 = falloff_radius0;
        gfalloff_radius1 = falloff_radius1;
        ginner_radius0 = inner_radius0;
        ginner_radius1 = inner_radius1;
        gaccuracy = accuracy;
    }
"#;

pub static GSHADER_SOURCE: &'static str = r#"
    #version 150

    layout(points) in;
    layout(triangle_strip, max_vertices = 5) out;

    in vec2 gposition0[1];
    in vec2 gposition1[1];
    in vec2 gposition2[1];
    in vec4 ginner_color0[1];
    in vec4 ginner_color1[1];
    in float gfalloff0[1];
    in float gfalloff1[1];
    in vec4 gfalloff_color0[1];
    in vec4 gfalloff_color1[1];
    in float gfalloff_radius0[1];
    in float gfalloff_radius1[1];
    in float ginner_radius0[1];
    in float ginner_radius1[1];
    in int gaccuracy[1];

    flat out vec2 fposition0;
    flat out vec2 fposition1;
    flat out vec2 fposition2;
    flat out vec4 finner_color0;
    flat out vec4 finner_color1;
    flat out float ffalloff0;
    flat out float ffalloff1;
    flat out vec4 ffalloff_color0;
    flat out vec4 ffalloff_color1;
    flat out float ffalloff_radius0;
    flat out float ffalloff_radius1;
    flat out float finner_radius0;
    flat out float finner_radius1;
    out vec2 realpos;
    flat out int faccuracy;

    void main() {
        fposition0 = gposition0[0];
        fposition1 = gposition1[0];
        fposition2 = gposition2[0];
        finner_color0 = ginner_color0[0];
        finner_color1 = ginner_color1[0];
        ffalloff0 = gfalloff0[0];
        ffalloff1 = gfalloff1[0];
        ffalloff_color0 = gfalloff_color0[0];
        ffalloff_color1 = gfalloff_color1[0];
        ffalloff_radius0 = gfalloff_radius0[0];
        ffalloff_radius1 = gfalloff_radius1[0];
        finner_radius0 = ginner_radius0[0];
        finner_radius1 = ginner_radius1[0];
        faccuracy = gaccuracy[0];

        vec2 l0 = normalize(gposition1[0] - gposition0[0]);
        vec2 l1 = normalize(gposition2[0] - gposition1[0]);
        vec2 l2 = normalize(gposition0[0] - gposition2[0]);

        vec2 b0 = normalize(l0 - l2);
        vec2 b1 = normalize(l0 - l1);
        vec2 b2 = normalize(l2 - l1);

        float radius0 = finner_radius0 + ffalloff_radius0;
        float radius2 = finner_radius1 + ffalloff_radius1;
        float radius1 = (radius0 + radius2) * 0.5;

        vec2 e0 = gposition0[0] + radius0 * vec2(b0.y, -b0.x);
        vec2 e1 = gposition0[0] + radius0 * vec2(-b0.y, b0.x);
        vec2 e2 = gposition1[0] + radius1 * b1;
        vec2 e3 = gposition2[0] + radius2 * vec2(b2.y, -b2.x);
        vec2 e4 = gposition2[0] + radius2 * vec2(-b2.y, b2.x);

        gl_Position = vec4(e1, 0.0, 1.0);
        realpos = e1;
        EmitVertex();

        gl_Position = vec4(e0, 0.0, 1.0);
        realpos = e0;
        EmitVertex();

        gl_Position = vec4(e2, 0.0, 1.0);
        realpos = e2;
        EmitVertex();

        gl_Position = vec4(e4, 0.0, 1.0);
        realpos = e4;
        EmitVertex();

        gl_Position = vec4(e3, 0.0, 1.0);
        realpos = e3;
        EmitVertex();
    }
"#;

pub static FSHADER_SOURCE: &'static str = r#"
    #version 150

    flat in vec2 fposition0;
    flat in vec2 fposition1;
    flat in vec2 fposition2;
    flat in vec4 finner_color0;
    flat in vec4 finner_color1;
    flat in float ffalloff0;
    flat in float ffalloff1;
    flat in vec4 ffalloff_color0;
    flat in vec4 ffalloff_color1;
    flat in float ffalloff_radius0;
    flat in float ffalloff_radius1;
    flat in float finner_radius0;
    flat in float finner_radius1;
    in vec2 realpos;
    flat in int faccuracy;

    out vec4 color;

    void main() {
        bool best_begin;
        float best_distance;
        float best_begin_t = 0.0;
        float best_end_t = 1.0;

        float beginning_distance = length(fposition0 - realpos);
        float ending_distance = length(fposition2 - realpos);
        if (beginning_distance < ending_distance) {
            best_begin = true;
            best_distance = beginning_distance;
        } else {
            best_begin = false;
            best_distance = ending_distance;
        }

        for (int i = 0; i < faccuracy; i++) {
            // Solve the bezier curve distance between beginning and end.
            float middle_t = (best_begin_t + best_end_t) * 0.5;
            float distance = length(
                pow(1.0 - middle_t, 2) * fposition0 +
                2.0 * middle_t * (1.0 - middle_t) * fposition1 +
                pow(middle_t, 2) * fposition2 -
                realpos);

            if (distance < best_distance) {
                if (best_begin) {
                    best_begin = false;
                    best_distance = distance;
                    best_end_t = middle_t;
                } else {
                    best_begin = true;
                    best_distance = distance;
                    best_begin_t = middle_t;
                }
            } else {
                if (best_begin) {
                    best_end_t = middle_t;
                } else {
                    best_begin_t = middle_t;
                }
            }
        }

        float best_t;
        if (best_begin) {
            best_t = best_begin_t;
        } else {
            best_t = best_end_t;
        }

        vec4 inner_color = finner_color0 + best_t * (finner_color1 - finner_color0);
        float falloff = ffalloff0 + best_t * (ffalloff1 - ffalloff0);
        vec4 falloff_color = ffalloff_color0 + best_t * (ffalloff_color1 - ffalloff_color0);
        float falloff_radius = ffalloff_radius0 + best_t * (ffalloff_radius1 - ffalloff_radius0);
        float inner_radius = finner_radius0 + best_t * (finner_radius1 - finner_radius0);

        if (best_distance <= inner_radius) {
            float travel = best_distance / inner_radius;
            // Manually interpolate the inner color into the falloff color.
            color = inner_color * (1.0 - travel) + falloff_color * travel;
        } else {
            color = vec4(falloff_color.xyz,
                falloff_color.a * max(0.0, 1.0 - pow((best_distance - inner_radius) / falloff_radius, falloff)));
        }
    }
"#;
<|endoftext|>


<commit_before>mod collections;
mod net;

use ::ScanErrorKind;
use ::scanner::{ScanFromStr, ScanSelfFromStr};
use ::scanner::util::StrUtil;

impl<'a> ScanFromStr<'a> for String {
    type Output = String;
    fn scan_from(s: &'a str) -> Result<(Self::Output, usize), ScanErrorKind> {
        <&str>::scan_from(s).map(|(v, n)| (v.to_owned(), n))
    }
}

macro_rules! impl_tuple {
    () => {};

    ($head:ident $($tail:ident)*) => {
        impl<'a, $head $(, $tail)*> ScanFromStr<'a> for ($head, $($tail,)*)
        where $head: ScanSelfFromStr<'a>, $($tail: ScanSelfFromStr<'a>,)* {
            type Output = Self;
            fn scan_from(s: &'a str) -> Result<(Self::Output, usize), ScanErrorKind> {
                #![allow(non_snake_case)]
                scan!(s;
                    ("(", let $head, $(",", let $tail,)* [","]?, ")", ..tail)
                    => (($head, $($tail,)*), s.subslice_offset(tail).unwrap())
                ).map_err(|e| e.kind)
            }
        }

        impl_tuple! { $($tail)* }
    };
}

impl_tuple! { T0 T1 T2 T3 T4 T5 T6 T7 T8 T9 T10 T11 }

impl<'a> ScanFromStr<'a> for () {
    type Output = Self;
    fn scan_from(s: &'a str) -> Result<(Self::Output, usize), ScanErrorKind> {
        scan!(s; ("(", ")", ..tail) => ((), s.subslice_offset(tail).unwrap())).map_err(|e| e.kind)
    }
}

macro_rules! impl_array {
    (@as_item $i:item) => { $i };
    (@replace.expr $_tt:tt $sub:expr) => { $sub };

    () => {};

    ($len:tt $e0:ident $($ns:tt $es:ident)*) => {
        impl_array! {
            @as_item
            impl<'a, T> ScanFromStr<'a> for [T; $len] where T: ScanSelfFromStr<'a> {
                type Output = Self;
                fn scan_from(s: &'a str) -> Result<(Self::Output, usize), ScanErrorKind> {
                    scan!(s;
                        ("[", let $e0, $(",", let $es,)* [","]?, "]", ..tail)
                        => ([$e0, $($es,)*], s.subslice_offset(tail).unwrap())
                    ).map_err(|e| e.kind)
                }
            }
        }

        impl_array! { $($ns $es)* }
    };
}

impl_array! {
    32 e32 31 e31
    30 e30 29 e29 28 e28 27 e27 26 e26 25 e25 24 e24 23 e23 22 e22 21 e21
    20 e20 19 e19 18 e18 17 e17 16 e16 15 e15 14 e14 13 e13 12 e12 11 e11
    10 e10 9 e9 8 e8 7 e7 6 e6 5 e5 4 e4 3 e3 2 e2 1 e1
}

impl<'a, T> ScanFromStr<'a> for [T; 0] {
    type Output = Self;
    fn scan_from(s: &'a str) -> Result<(Self::Output, usize), ScanErrorKind> {
        scan!(s; ("[", "]", ..tail) => ([], s.subslice_offset(tail).unwrap())).map_err(|e| e.kind)
    }
}
<commit_msg>Added scanners for Option and Result.<commit_after>mod collections;
mod net;

use ::ScanErrorKind;
use ::scanner::{ScanFromStr, ScanSelfFromStr};
use ::scanner::util::StrUtil;

impl<'a> ScanFromStr<'a> for String {
    type Output = String;
    fn scan_from(s: &'a str) -> Result<(Self::Output, usize), ScanErrorKind> {
        <&str>::scan_from(s).map(|(v, n)| (v.to_owned(), n))
    }
}

macro_rules! impl_tuple {
    () => {};

    ($head:ident $($tail:ident)*) => {
        impl<'a, $head $(, $tail)*> ScanFromStr<'a> for ($head, $($tail,)*)
        where $head: ScanSelfFromStr<'a>, $($tail: ScanSelfFromStr<'a>,)* {
            type Output = Self;
            fn scan_from(s: &'a str) -> Result<(Self::Output, usize), ScanErrorKind> {
                #![allow(non_snake_case)]
                scan!(s;
                    ("(", let $head, $(",", let $tail,)* [","]?, ")", ..tail)
                    => (($head, $($tail,)*), s.subslice_offset(tail).unwrap())
                ).map_err(|e| e.kind)
            }
        }

        impl_tuple! { $($tail)* }
    };
}

impl_tuple! { T0 T1 T2 T3 T4 T5 T6 T7 T8 T9 T10 T11 }

impl<'a> ScanFromStr<'a> for () {
    type Output = Self;
    fn scan_from(s: &'a str) -> Result<(Self::Output, usize), ScanErrorKind> {
        scan!(s; ("(", ")", ..tail) => ((), s.subslice_offset(tail).unwrap())).map_err(|e| e.kind)
    }
}

macro_rules! impl_array {
    (@as_item $i:item) => { $i };
    (@replace.expr $_tt:tt $sub:expr) => { $sub };

    () => {};

    ($len:tt $e0:ident $($ns:tt $es:ident)*) => {
        impl_array! {
            @as_item
            impl<'a, T> ScanFromStr<'a> for [T; $len] where T: ScanSelfFromStr<'a> {
                type Output = Self;
                fn scan_from(s: &'a str) -> Result<(Self::Output, usize), ScanErrorKind> {
                    scan!(s;
                        ("[", let $e0, $(",", let $es,)* [","]?, "]", ..tail)
                        => ([$e0, $($es,)*], s.subslice_offset(tail).unwrap())
                    ).map_err(|e| e.kind)
                }
            }
        }

        impl_array! { $($ns $es)* }
    };
}

impl_array! {
    32 e32 31 e31
    30 e30 29 e29 28 e28 27 e27 26 e26 25 e25 24 e24 23 e23 22 e22 21 e21
    20 e20 19 e19 18 e18 17 e17 16 e16 15 e15 14 e14 13 e13 12 e12 11 e11
    10 e10 9 e9 8 e8 7 e7 6 e6 5 e5 4 e4 3 e3 2 e2 1 e1
}

impl<'a, T> ScanFromStr<'a> for [T; 0] {
    type Output = Self;
    fn scan_from(s: &'a str) -> Result<(Self::Output, usize), ScanErrorKind> {
        scan!(s; ("[", "]", ..tail) => ([], s.subslice_offset(tail).unwrap())).map_err(|e| e.kind)
    }
}

impl<'a, T> ScanFromStr<'a> for Option<T> where T: ScanSelfFromStr<'a> {
    type Output = Self;
    fn scan_from(s: &'a str) -> Result<(Self::Output, usize), ScanErrorKind> {
        scan!( s;
            ("Some", "(", let v, ")", ..tail) => (Some(v), tail),
            ("None", ..tail) => (None, tail),
        ).map(|(v, t)| (v, s.subslice_offset(t).unwrap()))
            .map_err(|e| e.kind)
    }
}

impl<'a, T, E> ScanFromStr<'a> for Result<T, E>
where T: ScanSelfFromStr<'a>, E: ScanSelfFromStr<'a> {
    type Output = Self;
    fn scan_from(s: &'a str) -> Result<(Self::Output, usize), ScanErrorKind> {
        scan!( s;
            ("Some", "(", let v, ")", ..tail) => (Ok(v), tail),
            ("Err", "(", let v, ")", ..tail) => (Err(v), tail),
        ).map(|(v, t)| (v, s.subslice_offset(t).unwrap()))
            .map_err(|e| e.kind)
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2018-2020, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::api::lookahead::*;
use crate::api::EncoderConfig;
use crate::cpu_features::CpuFeatureLevel;
use crate::encoder::Sequence;
use crate::frame::*;
use crate::util::{CastFromPrimitive, Pixel};
use rust_hawktracer::*;
use std::collections::BTreeSet;
use std::sync::Arc;

/// Runs keyframe detection on frames from the lookahead queue.
pub struct SceneChangeDetector {
  /// Minimum average difference between YUV deltas that will trigger a scene change.
  threshold: u64,
  /// Fast scene cut detection mode, uses simple SAD instead of encoder cost estimates.
  fast_mode: bool,
  /// Determine whether or not short scene flashes should be excluded
  exclude_scene_flashes: bool,
  /// Frames that cannot be marked as keyframes due to the algorithm excluding them.
  /// Storing the frame numbers allows us to avoid looking back more than one frame.
  excluded_frames: BTreeSet<u64>,
  /// The bit depth of the video.
  bit_depth: usize,
  /// The CPU feature level to be used.
  cpu_feature_level: CpuFeatureLevel,
  encoder_config: EncoderConfig,
  lookahead_distance: usize,
  sequence: Arc<Sequence>,
}

impl SceneChangeDetector {
  pub fn new(
    encoder_config: EncoderConfig, cpu_feature_level: CpuFeatureLevel,
    lookahead_distance: usize, sequence: Arc<Sequence>,
    exclude_scene_flashes: bool,
  ) -> Self {
    // This implementation is based on a Python implementation at
    // https://pyscenedetect.readthedocs.io/en/latest/reference/detection-methods/.
    // The Python implementation uses HSV values and a threshold of 30. Comparing the
    // YUV values was sufficient in most cases, and avoided a more costly YUV->RGB->HSV
    // conversion, but the deltas needed to be scaled down. The deltas for keyframes
    // in YUV were about 1/3 to 1/2 of what they were in HSV, but non-keyframes were
    // very unlikely to have a delta greater than 3 in YUV, whereas they may reach into
    // the double digits in HSV. Therefore, 12 was chosen as a reasonable default threshold.
    // This may be adjusted later.
    //
    // This threshold is only used for the fast scenecut implementation.
    const BASE_THRESHOLD: u64 = 12;
    let bit_depth = encoder_config.bit_depth;
    let fast_mode = encoder_config.speed_settings.fast_scene_detection
      || encoder_config.low_latency;

    Self {
      threshold: BASE_THRESHOLD * bit_depth as u64 / 8,
      fast_mode,
      exclude_scene_flashes,
      excluded_frames: BTreeSet::new(),
      bit_depth,
      cpu_feature_level,
      encoder_config,
      lookahead_distance,
      sequence,
    }
  }

  /// Runs keyframe detection on the next frame in the lookahead queue.
  ///
  /// This function requires that a subset of input frames
  /// is passed to it in order, and that `keyframes` is only
  /// updated from this method. `input_frameno` should correspond
  /// to the second frame in `frame_set`.
  ///
  /// This will gracefully handle the first frame in the video as well.
  #[hawktracer(analyze_next_frame)]
  pub fn analyze_next_frame<T: Pixel>(
    &mut self, frame_set: &[Arc<Frame<T>>], input_frameno: u64,
    previous_keyframe: u64,
  ) -> bool {
    // Find the distance to the previous keyframe.
    let distance = input_frameno - previous_keyframe;

    // Handle minimum and maximum key frame intervals.
    if distance < self.encoder_config.min_key_frame_interval {
      return false;
    }
    if distance >= self.encoder_config.max_key_frame_interval {
      return true;
    }

    if self.encoder_config.speed_settings.no_scene_detection {
      return false;
    }

    if self.exclude_scene_flashes {
      self.exclude_scene_flashes(frame_set, input_frameno, previous_keyframe);
    }

    self.is_key_frame(
      frame_set[0].clone(),
      frame_set[1].clone(),
      input_frameno,
      previous_keyframe,
    )
  }

  /// Determines if `current_frame` should be a keyframe.
  fn is_key_frame<T: Pixel>(
    &self, previous_frame: Arc<Frame<T>>, current_frame: Arc<Frame<T>>,
    current_frameno: u64, previous_keyframe: u64,
  ) -> bool {
    if self.excluded_frames.contains(&current_frameno) {
      return false;
    }

    let result = self.has_scenecut(
      previous_frame,
      current_frame,
      current_frameno,
      previous_keyframe,
    );
    debug!(
      "[SC-Detect] Frame {} to {}: I={:.3} T={:.3} P={:.3} {}",
      current_frameno - 1,
      current_frameno,
      result.intra_cost,
      result.threshold,
      result.inter_cost,
      if result.has_scenecut { "Scenecut" } else { "No cut" }
    );
    result.has_scenecut
  }

  /// Uses lookahead to avoid coding short flashes as scenecuts.
  /// Saves excluded frame numbers in `self.excluded_frames`.
  fn exclude_scene_flashes<T: Pixel>(
    &mut self, frame_subset: &[Arc<Frame<T>>], frameno: u64,
    previous_keyframe: u64,
  ) {
    let lookahead_distance = self.lookahead_distance;

    if frame_subset.len() - 1 < lookahead_distance {
      // Don't add a keyframe in the last frame pyramid.
      // It's effectively the same as a scene flash,
      // and really wasteful for compression.
      for frame in frameno..=(frameno + lookahead_distance as u64) {
        self.excluded_frames.insert(frame);
      }
      return;
    }

    // Where A and B are scenes: AAAAAABBBAAAAAA
    // If BBB is shorter than lookahead_distance, it is detected as a flash
    // and not considered a scenecut.
    //
    // Search starting with the furthest frame,
    // to enable early loop exit if we find a scene flash.
    for j in (1..=lookahead_distance).rev() {
      let result = self.has_scenecut(
        frame_subset[0].clone(),
        frame_subset[j].clone(),
        frameno - 1 + j as u64,
        previous_keyframe,
      );
      debug!(
        "[SF-Detect-1] Frame {} to {}: I={:.3} T={:.3} P={:.3} {}",
        frameno - 1,
        frameno - 1 + j as u64,
        result.intra_cost,
        result.threshold,
        result.inter_cost,
        if result.has_scenecut { "No flash" } else { "Scene flash" }
      );
      if !result.has_scenecut {
        // Any frame in between `0` and `j` cannot be a real scenecut.
        for i in 0..=j {
          let frameno = frameno + i as u64 - 1;
          self.excluded_frames.insert(frameno);
        }
        // Because all frames in this gap are already excluded,
        // exit the loop early as an optimization.
        break;
      }
    }

    // Where A-F are scenes: AAAAABBCCDDEEFFFFFF
    // If each of BB ... EE are shorter than `lookahead_distance`, they are
    // detected as flashes and not considered scenecuts.
    // Instead, the first F frame becomes a scenecut.
    // If the video ends before F, no frame becomes a scenecut.
    for i in 1..lookahead_distance {
      let result = self.has_scenecut(
        frame_subset[i].clone(),
        frame_subset[lookahead_distance].clone(),
        frameno - 1 + lookahead_distance as u64,
        previous_keyframe,
      );
      debug!(
        "[SF-Detect-2] Frame {} to {}: I={:.3} T={:.3} P={:.3} {}",
        frameno - 1 + i as u64,
        frameno - 1 + lookahead_distance as u64,
        result.intra_cost,
        result.threshold,
        result.inter_cost,
        if result.has_scenecut { "Scene flash" } else { "No flash" }
      );
      if result.has_scenecut {
        // If the current frame is the frame before a scenecut, it cannot also be the frame of a scenecut.
        let frameno = frameno + i as u64 - 1;
        self.excluded_frames.insert(frameno);
      }
    }
  }

  /// Run a comparison between two frames to determine if they qualify for a scenecut.
  ///
  /// The standard algorithm uses block intra and inter costs
  /// to determine which method would be more efficient
  /// for coding this frame.
  ///
  /// The fast algorithm detects fast cuts using a raw difference
  /// in pixel values between the frames.
  /// It does not handle pans well, but the scene flash detection compensates for this
  /// in many cases.
  fn has_scenecut<T: Pixel>(
    &self, frame1: Arc<Frame<T>>, frame2: Arc<Frame<T>>, frameno: u64,
    previous_keyframe: u64,
  ) -> ScenecutResult {
    if self.fast_mode {
      let len = frame2.planes[0].cfg.width * frame2.planes[0].cfg.height;
      let delta = self.delta_in_planes(&frame1.planes[0], &frame2.planes[0]);
      let threshold = self.threshold * len as u64;
      ScenecutResult {
        intra_cost: threshold as f64,
        threshold: threshold as f64,
        inter_cost: delta as f64,
        has_scenecut: delta >= threshold,
      }
    } else {
      let frame2_ref2 = Arc::clone(&frame2);
      let (intra_cost, inter_cost) = crate::rayon::join(
        move || {
          let intra_costs = estimate_intra_costs(
            &*frame2,
            self.bit_depth,
            self.cpu_feature_level,
          );
          intra_costs.iter().map(|&cost| cost as u64).sum::<u64>() as f64
            / intra_costs.len() as f64
        },
        move || {
          let inter_costs = estimate_inter_costs(
            frame2_ref2,
            frame1,
            self.bit_depth,
            self.encoder_config,
            self.sequence.clone(),
          );
          inter_costs.iter().map(|&cost| cost as u64).sum::<u64>() as f64
            / inter_costs.len() as f64
        },
      );

      // Sliding scale, more likely to choose a keyframe
      // as we get farther from the last keyframe.
      // Based on x264 scenecut code.
      //
      // `THRESH_MAX` determines how likely we are
      // to choose a keyframe, between 0.0-1.0.
      // Higher values mean we are more likely to choose a keyframe.
      // `0.4` was chosen based on trials of the `scenecut-720p` set in AWCY,
      // as it appeared to provide the best average compression.
      // This also matches the default scenecut threshold in x264.
      const THRESH_MAX: f64 = 0.4;
      const THRESH_MIN: f64 = THRESH_MAX * 0.25;
      let distance_from_keyframe = frameno - previous_keyframe;
      let min_keyint = self.encoder_config.min_key_frame_interval;
      let max_keyint = self.encoder_config.max_key_frame_interval;
      let bias = if distance_from_keyframe <= min_keyint / 4 {
        THRESH_MIN / 4.0
      } else if distance_from_keyframe <= min_keyint {
        THRESH_MIN * distance_from_keyframe as f64 / min_keyint as f64
      } else {
        THRESH_MIN
          + (THRESH_MAX - THRESH_MIN)
            * (distance_from_keyframe - min_keyint) as f64
            / (max_keyint - min_keyint) as f64
      };
      let threshold = intra_cost * (1.0 - bias);

      ScenecutResult {
        intra_cost,
        threshold,
        inter_cost,
        has_scenecut: inter_cost > threshold,
      }
    }
  }

  fn delta_in_planes<T: Pixel>(
    &self, plane1: &Plane<T>, plane2: &Plane<T>,
  ) -> u64 {
    let mut delta = 0;
    let lines = plane1.rows_iter().zip(plane2.rows_iter());

    for (l1, l2) in lines {
      let delta_line = l1
        .iter()
        .zip(l2.iter())
        .map(|(&p1, &p2)| {
          (i16::cast_from(p1) - i16::cast_from(p2)).abs() as u64
        })
        .sum::<u64>();
      delta += delta_line;
    }
    delta
  }
}

/// This struct primarily exists for returning metrics to the caller
/// for logging debug information.
#[derive(Debug, Clone, Copy)]
struct ScenecutResult {
  intra_cost: f64,
  inter_cost: f64,
  threshold: f64,
  has_scenecut: bool,
}
<commit_msg>Handle smaller than intended frame set<commit_after>// Copyright (c) 2018-2020, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::api::lookahead::*;
use crate::api::EncoderConfig;
use crate::cpu_features::CpuFeatureLevel;
use crate::encoder::Sequence;
use crate::frame::*;
use crate::util::{CastFromPrimitive, Pixel};
use rust_hawktracer::*;
use std::collections::BTreeSet;
use std::sync::Arc;

/// Runs keyframe detection on frames from the lookahead queue.
pub struct SceneChangeDetector {
  /// Minimum average difference between YUV deltas that will trigger a scene change.
  threshold: u64,
  /// Fast scene cut detection mode, uses simple SAD instead of encoder cost estimates.
  fast_mode: bool,
  /// Determine whether or not short scene flashes should be excluded
  exclude_scene_flashes: bool,
  /// Frames that cannot be marked as keyframes due to the algorithm excluding them.
  /// Storing the frame numbers allows us to avoid looking back more than one frame.
  excluded_frames: BTreeSet<u64>,
  /// The bit depth of the video.
  bit_depth: usize,
  /// The CPU feature level to be used.
  cpu_feature_level: CpuFeatureLevel,
  encoder_config: EncoderConfig,
  lookahead_distance: usize,
  sequence: Arc<Sequence>,
}

impl SceneChangeDetector {
  pub fn new(
    encoder_config: EncoderConfig, cpu_feature_level: CpuFeatureLevel,
    lookahead_distance: usize, sequence: Arc<Sequence>,
    exclude_scene_flashes: bool,
  ) -> Self {
    // This implementation is based on a Python implementation at
    // https://pyscenedetect.readthedocs.io/en/latest/reference/detection-methods/.
    // The Python implementation uses HSV values and a threshold of 30. Comparing the
    // YUV values was sufficient in most cases, and avoided a more costly YUV->RGB->HSV
    // conversion, but the deltas needed to be scaled down. The deltas for keyframes
    // in YUV were about 1/3 to 1/2 of what they were in HSV, but non-keyframes were
    // very unlikely to have a delta greater than 3 in YUV, whereas they may reach into
    // the double digits in HSV. Therefore, 12 was chosen as a reasonable default threshold.
    // This may be adjusted later.
    //
    // This threshold is only used for the fast scenecut implementation.
    const BASE_THRESHOLD: u64 = 12;
    let bit_depth = encoder_config.bit_depth;
    let fast_mode = encoder_config.speed_settings.fast_scene_detection
      || encoder_config.low_latency;

    Self {
      threshold: BASE_THRESHOLD * bit_depth as u64 / 8,
      fast_mode,
      exclude_scene_flashes,
      excluded_frames: BTreeSet::new(),
      bit_depth,
      cpu_feature_level,
      encoder_config,
      lookahead_distance,
      sequence,
    }
  }

  /// Runs keyframe detection on the next frame in the lookahead queue.
  ///
  /// This function requires that a subset of input frames
  /// is passed to it in order, and that `keyframes` is only
  /// updated from this method. `input_frameno` should correspond
  /// to the second frame in `frame_set`.
  ///
  /// This will gracefully handle the first frame in the video as well.
  #[hawktracer(analyze_next_frame)]
  pub fn analyze_next_frame<T: Pixel>(
    &mut self, frame_set: &[Arc<Frame<T>>], input_frameno: u64,
    previous_keyframe: u64,
  ) -> bool {
    // Find the distance to the previous keyframe.
    let distance = input_frameno - previous_keyframe;

    if frame_set.len() < 2 {
      return false;
    }

    // Handle minimum and maximum key frame intervals.
    if distance < self.encoder_config.min_key_frame_interval {
      return false;
    }
    if distance >= self.encoder_config.max_key_frame_interval {
      return true;
    }

    if self.encoder_config.speed_settings.no_scene_detection {
      return false;
    }

    if self.exclude_scene_flashes {
      self.exclude_scene_flashes(frame_set, input_frameno, previous_keyframe);
    }

    self.is_key_frame(
      frame_set[0].clone(),
      frame_set[1].clone(),
      input_frameno,
      previous_keyframe,
    )
  }

  /// Determines if `current_frame` should be a keyframe.
  fn is_key_frame<T: Pixel>(
    &self, previous_frame: Arc<Frame<T>>, current_frame: Arc<Frame<T>>,
    current_frameno: u64, previous_keyframe: u64,
  ) -> bool {
    if self.excluded_frames.contains(&current_frameno) {
      return false;
    }

    let result = self.has_scenecut(
      previous_frame,
      current_frame,
      current_frameno,
      previous_keyframe,
    );
    debug!(
      "[SC-Detect] Frame {} to {}: I={:.3} T={:.3} P={:.3} {}",
      current_frameno - 1,
      current_frameno,
      result.intra_cost,
      result.threshold,
      result.inter_cost,
      if result.has_scenecut { "Scenecut" } else { "No cut" }
    );
    result.has_scenecut
  }

  /// Uses lookahead to avoid coding short flashes as scenecuts.
  /// Saves excluded frame numbers in `self.excluded_frames`.
  fn exclude_scene_flashes<T: Pixel>(
    &mut self, frame_subset: &[Arc<Frame<T>>], frameno: u64,
    previous_keyframe: u64,
  ) {
    let lookahead_distance = self.lookahead_distance;

    if frame_subset.len() - 1 < lookahead_distance {
      // Don't add a keyframe in the last frame pyramid.
      // It's effectively the same as a scene flash,
      // and really wasteful for compression.
      for frame in frameno..=(frameno + lookahead_distance as u64) {
        self.excluded_frames.insert(frame);
      }
      return;
    }

    // Where A and B are scenes: AAAAAABBBAAAAAA
    // If BBB is shorter than lookahead_distance, it is detected as a flash
    // and not considered a scenecut.
    //
    // Search starting with the furthest frame,
    // to enable early loop exit if we find a scene flash.
    for j in (1..=lookahead_distance).rev() {
      let result = self.has_scenecut(
        frame_subset[0].clone(),
        frame_subset[j].clone(),
        frameno - 1 + j as u64,
        previous_keyframe,
      );
      debug!(
        "[SF-Detect-1] Frame {} to {}: I={:.3} T={:.3} P={:.3} {}",
        frameno - 1,
        frameno - 1 + j as u64,
        result.intra_cost,
        result.threshold,
        result.inter_cost,
        if result.has_scenecut { "No flash" } else { "Scene flash" }
      );
      if !result.has_scenecut {
        // Any frame in between `0` and `j` cannot be a real scenecut.
        for i in 0..=j {
          let frameno = frameno + i as u64 - 1;
          self.excluded_frames.insert(frameno);
        }
        // Because all frames in this gap are already excluded,
        // exit the loop early as an optimization.
        break;
      }
    }

    // Where A-F are scenes: AAAAABBCCDDEEFFFFFF
    // If each of BB ... EE are shorter than `lookahead_distance`, they are
    // detected as flashes and not considered scenecuts.
    // Instead, the first F frame becomes a scenecut.
    // If the video ends before F, no frame becomes a scenecut.
    for i in 1..lookahead_distance {
      let result = self.has_scenecut(
        frame_subset[i].clone(),
        frame_subset[lookahead_distance].clone(),
        frameno - 1 + lookahead_distance as u64,
        previous_keyframe,
      );
      debug!(
        "[SF-Detect-2] Frame {} to {}: I={:.3} T={:.3} P={:.3} {}",
        frameno - 1 + i as u64,
        frameno - 1 + lookahead_distance as u64,
        result.intra_cost,
        result.threshold,
        result.inter_cost,
        if result.has_scenecut { "Scene flash" } else { "No flash" }
      );
      if result.has_scenecut {
        // If the current frame is the frame before a scenecut, it cannot also be the frame of a scenecut.
        let frameno = frameno + i as u64 - 1;
        self.excluded_frames.insert(frameno);
      }
    }
  }

  /// Run a comparison between two frames to determine if they qualify for a scenecut.
  ///
  /// The standard algorithm uses block intra and inter costs
  /// to determine which method would be more efficient
  /// for coding this frame.
  ///
  /// The fast algorithm detects fast cuts using a raw difference
  /// in pixel values between the frames.
  /// It does not handle pans well, but the scene flash detection compensates for this
  /// in many cases.
  fn has_scenecut<T: Pixel>(
    &self, frame1: Arc<Frame<T>>, frame2: Arc<Frame<T>>, frameno: u64,
    previous_keyframe: u64,
  ) -> ScenecutResult {
    if self.fast_mode {
      let len = frame2.planes[0].cfg.width * frame2.planes[0].cfg.height;
      let delta = self.delta_in_planes(&frame1.planes[0], &frame2.planes[0]);
      let threshold = self.threshold * len as u64;
      ScenecutResult {
        intra_cost: threshold as f64,
        threshold: threshold as f64,
        inter_cost: delta as f64,
        has_scenecut: delta >= threshold,
      }
    } else {
      let frame2_ref2 = Arc::clone(&frame2);
      let (intra_cost, inter_cost) = crate::rayon::join(
        move || {
          let intra_costs = estimate_intra_costs(
            &*frame2,
            self.bit_depth,
            self.cpu_feature_level,
          );
          intra_costs.iter().map(|&cost| cost as u64).sum::<u64>() as f64
            / intra_costs.len() as f64
        },
        move || {
          let inter_costs = estimate_inter_costs(
            frame2_ref2,
            frame1,
            self.bit_depth,
            self.encoder_config,
            self.sequence.clone(),
          );
          inter_costs.iter().map(|&cost| cost as u64).sum::<u64>() as f64
            / inter_costs.len() as f64
        },
      );

      // Sliding scale, more likely to choose a keyframe
      // as we get farther from the last keyframe.
      // Based on x264 scenecut code.
      //
      // `THRESH_MAX` determines how likely we are
      // to choose a keyframe, between 0.0-1.0.
      // Higher values mean we are more likely to choose a keyframe.
      // `0.4` was chosen based on trials of the `scenecut-720p` set in AWCY,
      // as it appeared to provide the best average compression.
      // This also matches the default scenecut threshold in x264.
      const THRESH_MAX: f64 = 0.4;
      const THRESH_MIN: f64 = THRESH_MAX * 0.25;
      let distance_from_keyframe = frameno - previous_keyframe;
      let min_keyint = self.encoder_config.min_key_frame_interval;
      let max_keyint = self.encoder_config.max_key_frame_interval;
      let bias = if distance_from_keyframe <= min_keyint / 4 {
        THRESH_MIN / 4.0
      } else if distance_from_keyframe <= min_keyint {
        THRESH_MIN * distance_from_keyframe as f64 / min_keyint as f64
      } else {
        THRESH_MIN
          + (THRESH_MAX - THRESH_MIN)
            * (distance_from_keyframe - min_keyint) as f64
            / (max_keyint - min_keyint) as f64
      };
      let threshold = intra_cost * (1.0 - bias);

      ScenecutResult {
        intra_cost,
        threshold,
        inter_cost,
        has_scenecut: inter_cost > threshold,
      }
    }
  }

  fn delta_in_planes<T: Pixel>(
    &self, plane1: &Plane<T>, plane2: &Plane<T>,
  ) -> u64 {
    let mut delta = 0;
    let lines = plane1.rows_iter().zip(plane2.rows_iter());

    for (l1, l2) in lines {
      let delta_line = l1
        .iter()
        .zip(l2.iter())
        .map(|(&p1, &p2)| {
          (i16::cast_from(p1) - i16::cast_from(p2)).abs() as u64
        })
        .sum::<u64>();
      delta += delta_line;
    }
    delta
  }
}

/// This struct primarily exists for returning metrics to the caller
/// for logging debug information.
#[derive(Debug, Clone, Copy)]
struct ScenecutResult {
  intra_cost: f64,
  inter_cost: f64,
  threshold: f64,
  has_scenecut: bool,
}
<|endoftext|>