1 files changed, 13 insertions, 2 deletions
diff --git a/rst_parser/rst.h b/rst_parser/rst.h
index 865871eb..8bf389f7 100644
--- a/rst_parser/rst.h
+++ b/rst_parser/rst.h
@@ -1,10 +1,21 @@
 #ifndef _RST_H_
 #define _RST_H_
 
+#include <vector>
+#include "sampler.h"
 #include "arc_factored.h"
+#include "alias_sampler.h"
 
-struct StochasticForest {
-  explicit StochasticForest(const ArcFactoredForest& af);
+struct TreeSampler {
+  explicit TreeSampler(const ArcFactoredForest& af);
+  void SampleRandomSpanningTree(EdgeSubset* tree, MT19937* rng);
+  const ArcFactoredForest& forest;
+#define USE_ALIAS_SAMPLER 1
+#if USE_ALIAS_SAMPLER
+  std::vector<AliasSampler> usucc;
+#else
+  std::vector<SampleSet<double> > usucc;
+#endif
 };
 
 #endif