1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

use quote::format_ident;
use syn::Ident;

use crate::qs::{quote, Index, TokenStream};
use crate::FastIndexMap;

#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
pub enum BindResourceType {
  Buffer,
  Sampler,
  Texture,
}

#[derive(Clone)]
pub struct BindingGenerator {
  pub bind_group_layout: BindGroupLayoutGenerator,
  pub pipeline_layout: PipelineLayoutGenerator,
}

impl std::fmt::Debug for BindingGenerator {
  fn fmt(&self, _: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    // skip the debug generation for this,
    // as the output changes on every build due to fns
    Ok(())
  }
}

/// Represents a generator for creating WGSL bind group layout structures.
///
/// This struct is used to generate the code for creating a bind group layout in WGSL.
/// It contains the necessary information for generating the code, such as the prefix name for the layout,
/// whether the generated code uses lifetimes, the type of the entry struct, a function for constructing entries,
/// and a map of binding resource types to their corresponding token streams.
#[derive(Clone, Debug)]
pub struct BindGroupLayoutGenerator {
  /// The prefix for the bind group layout.
  pub name_prefix: String,

  /// Indicates whether the generated code uses lifetimes.
  ///
  /// If this is `true`, the generated code will include lifetimes. If it's `false`, the generated code will not include lifetimes.
  pub uses_lifetime: bool,

  /// The type of the entry struct in the generated code.
  ///
  /// This is represented as a `TokenStream` that contains the code for the type of the entry struct.
  pub entry_struct_type: TokenStream,

  /// A function for constructing entries in the generated code.
  ///
  /// This function takes a binding index, a `TokenStream` for the binding variable, and a `WgslBindResourceType` for the resource type,
  /// and returns a `TokenStream` that contains the code for constructing an entry.
  pub entry_constructor: fn(usize, TokenStream, BindResourceType) -> TokenStream,

  /// A map of binding resource types to their corresponding token streams.
  ///
  /// This map is used to generate the code for the binding resources in the bind group layout.
  pub binding_type_map: FastIndexMap<BindResourceType, TokenStream>,
}

impl BindGroupLayoutGenerator {
  pub(crate) fn bind_group_name_ident(&self, group_index: u32) -> Ident {
    format_ident!("{}{}", self.name_prefix, group_index)
  }

  pub(crate) fn bind_group_entry_collection_struct_name_ident(
    &self,
    group_index: u32,
  ) -> Ident {
    format_ident!("{}{}{}", self.name_prefix, group_index, "EntryCollection")
  }
}

#[derive(Clone, Debug)]
pub struct PipelineLayoutGenerator {
  pub layout_name: String,
  pub bind_group_layout_type: TokenStream,
}

pub trait GetBindingsGeneratorConfig {
  fn get_generator_config(self) -> BindingGenerator;
}
impl GetBindingsGeneratorConfig for BindingGenerator {
  fn get_generator_config(self) -> BindingGenerator {
    self
  }
}

impl Default for BindingGenerator {
  fn default() -> BindingGenerator {
    WgpuGetBindingsGeneratorConfig.get_generator_config()
  }
}

pub struct WgpuGetBindingsGeneratorConfig;
impl WgpuGetBindingsGeneratorConfig {
  fn get_bind_group_layout_generator_config() -> BindGroupLayoutGenerator {
    let binding_type_map = vec![
      (BindResourceType::Buffer, quote! { wgpu::BufferBinding<'a> }),
      (BindResourceType::Sampler, quote! { &'a wgpu::Sampler }),
      (BindResourceType::Texture, quote! { &'a wgpu::TextureView }),
    ]
    .into_iter()
    .collect::<FastIndexMap<_, _>>();

    fn entry_constructor(
      binding: usize,
      binding_var: TokenStream,
      resource_type: BindResourceType,
    ) -> TokenStream {
      let resource = match resource_type {
        BindResourceType::Buffer => {
          quote!(wgpu::BindingResource::Buffer(#binding_var))
        }
        BindResourceType::Sampler => {
          quote!(wgpu::BindingResource::Sampler(#binding_var))
        }
        BindResourceType::Texture => {
          quote!(wgpu::BindingResource::TextureView(#binding_var))
        }
      };

      let binding = Index::from(binding);
      quote! {
        wgpu::BindGroupEntry {
          binding: #binding,
          resource: #resource,
        }
      }
    }

    BindGroupLayoutGenerator {
      name_prefix: "WgpuBindGroup".into(),
      uses_lifetime: true,
      entry_struct_type: quote!(wgpu::BindGroupEntry<'a>),
      entry_constructor,
      binding_type_map,
    }
  }

  fn get_pipeline_layout_generator() -> PipelineLayoutGenerator {
    PipelineLayoutGenerator {
      layout_name: "WgpuPipelineLayout".into(),
      bind_group_layout_type: quote!(wgpu::BindGroupLayout),
    }
  }
}

impl GetBindingsGeneratorConfig for WgpuGetBindingsGeneratorConfig {
  fn get_generator_config(self) -> BindingGenerator {
    let bind_group_layout = Self::get_bind_group_layout_generator_config();
    let pipeline_layout = Self::get_pipeline_layout_generator();
    BindingGenerator {
      bind_group_layout,
      pipeline_layout,
    }
  }
}