add source type for context (#120)

cloudwego · Mar 17, 2023 · cf0af28 · cf0af28
1 parent 5aa7b4e
commit cf0af28
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 5 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/pilota-build/Cargo.toml b/pilota-build/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pilota-build"
-version = "0.6.0"
+version = "0.6.1"
 edition = "2021"
 description = "Compile thrift and protobuf idl into rust code at compile-time."
 documentation = "https://docs.rs/pilota-build"

diff --git a/pilota-build/src/lib.rs b/pilota-build/src/lib.rs
@@ -31,7 +31,7 @@ use db::{RirDatabase, RootDatabase};
 use fmt::fmt_file;
 pub use middle::{context::Context, rir, ty};
 use middle::{
-    context::{tls::CONTEXT, CollectMode},
+    context::{tls::CONTEXT, CollectMode, SourceType},
     rir::NodeKind,
     type_graph::TypeGraph,
 };
@@ -73,6 +73,7 @@ impl MakeBackend for MkProtobufBackend {
 }
 
 pub struct Builder<MkB, P> {
+    source_type: SourceType,
     mk_backend: MkB,
     parser: P,
     plugins: Vec<Box<dyn Plugin>>,
@@ -83,6 +84,7 @@ pub struct Builder<MkB, P> {
 impl Builder<MkThriftBackend, ThriftParser> {
     pub fn thrift() -> Self {
         Builder {
+            source_type: SourceType::Thrift,
             mk_backend: MkThriftBackend,
             parser: ThriftParser::default(),
             plugins: vec![
@@ -99,6 +101,7 @@ impl Builder<MkThriftBackend, ThriftParser> {
 impl Builder<MkProtobufBackend, ProtobufParser> {
     pub fn protobuf() -> Self {
         Builder {
+            source_type: SourceType::Protobuf,
             mk_backend: MkProtobufBackend,
             parser: ProtobufParser::default(),
             plugins: vec![
@@ -125,6 +128,7 @@ where
 impl<MkB, P> Builder<MkB, P> {
     pub fn with_backend<B: MakeBackend>(self, mk_backend: B) -> Builder<B, P> {
         Builder {
+            source_type: self.source_type,
             mk_backend,
             parser: self.parser,
             plugins: self.plugins,
@@ -195,7 +199,7 @@ where
         db.set_type_graph_with_durability(type_graph, Durability::HIGH);
         db.set_nodes_with_durability(Arc::new(nodes), Durability::HIGH);
 
-        let mut cx = Context::new(db.snapshot());
+        let mut cx = Context::new(self.source_type, db.snapshot());
         cx.set_tags_map(tags);
 
         cx.exec_plugin(BoxedPlugin);

diff --git a/pilota-build/src/middle/context.rs b/pilota-build/src/middle/context.rs
@@ -28,6 +28,7 @@ pub enum CollectMode {
 }
 
 pub struct Context {
+    pub source_type: SourceType,
     pub db: salsa::Snapshot<RootDatabase>,
     adjusts: FxHashMap<DefId, Adjust>,
     tags_map: FxHashMap<TagId, Arc<Tags>>,
@@ -41,9 +42,15 @@ impl Deref for Context {
     }
 }
 
+pub enum SourceType {
+    Thrift,
+    Protobuf,
+}
+
 impl Context {
-    pub fn new(db: salsa::Snapshot<RootDatabase>) -> Context {
+    pub fn new(source_type: SourceType, db: salsa::Snapshot<RootDatabase>) -> Context {
         Context {
+            source_type,
             db,
             adjusts: Default::default(),
             tags_map: Default::default(),