Adjust setRNGState so that it updates the pointer to the generator.

d11 · d11 · commit 0003434c29a7 · 2015-10-27T12:45:05.000Z
Otherwise, restoring a state from a previous run can set the pointer to
one that is no longer valid.
diff --git a/luasrc/tests/testWrap.lua b/luasrc/tests/testWrap.lua
@@ -54,6 +54,21 @@ function myTest.callsetRNGWithGenerator()
     tester:assert(ok, 'Failed to set RNG state')
 end
 
+function myTest.setRNGWithBadPointer()
+  -- To simulate restoring the state from a previous run, we invalidate the
+  -- pointer to the main Torch generator in a state that we are passing to
+  -- setRNGState, and check that this doesn't break things.
+  local state = torch.getRNGState()
+  local x = tonumber(randomkit.binomial(10, 0.4))
+  local badState = ffi.cast("rk_state *", state.randomkit)
+  badState.torch_state = ffi.cast("THGenerator*", 0)
+  torch.setRNGState{
+      torch = state.torch,
+      randomkit = badState
+  }
+  tester:asserteq(x, tonumber(randomkit.binomial(10, 0.4)))
+end
+
 
 tester:add(myTest)
 return tester:run()
diff --git a/luasrc/wrapC.lua b/luasrc/wrapC.lua
@@ -210,6 +210,12 @@ torch.setRNGState = function(generator, state)
     _setRNGState(state.torch)
     -- Deserialize from string
     ffi.copy(randomkit._state, state.randomkit, ffi.sizeof(randomkit._state))
+
+    -- If the state being set is from a previous run, the pointer to the main
+    -- Torch generator will no longer be valid. So we will explicitly update the
+    -- pointer to be current.
+    randomkit._state.torch_state =
+        ffi.cast("THGenerator*", torch.pointer(torch._gen))
 end
 
 local returnTypeMapping = {
