From acebaa517b5aed232040b092ef37708b3563199d Mon Sep 17 00:00:00 2001
From: Ozzie Gooen <ozzieagooen@gmail.com>
Date: Thu, 21 Jul 2022 10:31:53 -0700
Subject: [PATCH] Simple SampleSet.fromDist function

---
 ...leLibrary_FunctionRegistryLibrary_test.res |  1 +
 .../FunctionRegistry/Library/FR_Sampleset.res | 35 +++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/packages/squiggle-lang/__tests__/SquiggleLibrary/SquiggleLibrary_FunctionRegistryLibrary_test.res b/packages/squiggle-lang/__tests__/SquiggleLibrary/SquiggleLibrary_FunctionRegistryLibrary_test.res
index da357653..a7e98163 100644
--- a/packages/squiggle-lang/__tests__/SquiggleLibrary/SquiggleLibrary_FunctionRegistryLibrary_test.res
+++ b/packages/squiggle-lang/__tests__/SquiggleLibrary/SquiggleLibrary_FunctionRegistryLibrary_test.res
@@ -62,6 +62,7 @@ describe("FunctionRegistry Library", () => {
     testEvalToBe("Dist.klDivergence(normal(5,2), normal(5,1.5))", "Ok(0.06874342818671068)")
     testEvalToBe("SampleSet.fromList([3,5,2,3,5,2,3,5,2,3,3,5])", "Ok(Sample Set Distribution)")
     testEvalToBe("SampleSet.fromList([3,5,2,3,5,2,3,5,2,3,3,5])", "Ok(Sample Set Distribution)")
+    testEvalToBe("SampleSet.fromFn({|| sample(normal(5,2))})", "Ok(Sample Set Distribution)")
     testEvalToBe(
       "addOne(t)=t+1; SampleSet.toList(SampleSet.map(SampleSet.fromList([1,2,3,4,5,6]), addOne))",
       "Ok([2,3,4,5,6,7])",
diff --git a/packages/squiggle-lang/src/rescript/FunctionRegistry/Library/FR_Sampleset.res b/packages/squiggle-lang/src/rescript/FunctionRegistry/Library/FR_Sampleset.res
index a13ff46c..b2927b82 100644
--- a/packages/squiggle-lang/src/rescript/FunctionRegistry/Library/FR_Sampleset.res
+++ b/packages/squiggle-lang/src/rescript/FunctionRegistry/Library/FR_Sampleset.res
@@ -21,6 +21,17 @@ module Internal = {
     | Error(r) => Error(REDistributionError(SampleSetError(r)))
     }
 
+  //TODO: I don't know why this seems to need at least one input
+  let fromFn = (
+    aLambdaValue,
+    env: ReducerInterface_InternalExpressionValue.environment,
+    reducer,
+  ) => {
+    let sampleCount = env.sampleCount
+    let fn = (r) => doLambdaCall(aLambdaValue, list{IEvNumber(r)}, env, reducer)
+    Belt_Array.makeBy(sampleCount, (r) => fn(r->Js.Int.toFloat))->E.A.R.firstErrorOrOpen
+  }
+
   let map1 = (sampleSetDist: t, aLambdaValue, env, reducer) => {
     let fn = r => doLambdaCall(aLambdaValue, list{IEvNumber(r)}, env, reducer)
     SampleSetDist.samplesMap(~fn, sampleSetDist)->toType
@@ -132,6 +143,30 @@ let library = [
     ],
     (),
   ),
+  Function.make(
+    ~name="fromFn",
+    ~nameSpace,
+    ~requiresNamespace=false,
+    ~examples=[`SampleSet.fromFn(sample(normal(5,2)))`],
+    ~output=ReducerInterface_InternalExpressionValue.EvtArray,
+    ~definitions=[
+      FnDefinition.make(
+        ~name="fromFn",
+        ~inputs=[FRTypeLambda],
+        ~run=(inputs, _, env, reducer) =>
+          switch inputs {
+          | [IEvLambda(lambda)] =>
+            switch Internal.fromFn(lambda, env, reducer) {
+            | Ok(r) => Ok(r->Wrappers.sampleSet->Wrappers.evDistribution)
+            | Error(_) => Error("issue")
+            }
+          | _ => Error(impossibleError)
+          },
+        (),
+      ),
+    ],
+    (),
+  ),
   Function.make(
     ~name="map",
     ~nameSpace,